experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Make overly-specific ShapeUtil predicate a little more general.

Browse Source 
PiperOrigin-RevId: 216263039
This commit is contained in:
Chris Leary 2018-10-08 15:42:17 -07:00 committed by TensorFlower Gardener
parent eb0f862ba6
commit cb057ea640
4 changed files with 10 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
tensorflow/compiler/xla/service/hlo_instruction_test.cc
View File

@ -135,7 +135,8 @@ TEST_F(HloInstructionTest, BasicProperties) {
auto parameter = HloInstruction::CreateParameter(1, r0f32_, "foo");
EXPECT_EQ(HloOpcode::kParameter, parameter->opcode());
EXPECT_TRUE(ShapeUtil::IsScalarF32(parameter->shape()));
EXPECT_TRUE(ShapeUtil::IsScalarWithElementType(parameter->shape(), F32));
EXPECT_FALSE(ShapeUtil::IsScalarWithElementType(parameter->shape(), S32));
EXPECT_EQ(0, parameter->operand_count());
}

2
tensorflow/compiler/xla/service/hlo_query.cc
View File

@ -24,7 +24,7 @@ namespace hlo_query {
bool IsConstantR0F32(HloInstruction* instruction, float* out) {
if (instruction->opcode() == HloOpcode::kConstant &&
ShapeUtil::IsScalarF32(instruction->shape())) {
ShapeUtil::IsScalarWithElementType(instruction->shape(), F32)) {
*out = instruction->literal().Get<float>({});
return true;
}

5
tensorflow/compiler/xla/shape_util.cc
View File

@ -461,8 +461,9 @@ ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(
return ShapeUtil::IsArray(shape) && ElementsIn(shape) == 0;
}
/* static */ bool ShapeUtil::IsScalarF32(const Shape& shape) {
return shape.element_type() == F32 && Rank(shape) == 0;
/* static */ bool ShapeUtil::IsScalarWithElementType(
const Shape& shape, PrimitiveType element_type) {
return IsScalar(shape) && shape.element_type() == element_type;
}
namespace {

5
tensorflow/compiler/xla/shape_util.h
View File

@ -312,7 +312,10 @@ class ShapeUtil {
static bool IsEffectiveScalar(const Shape& shape) {
return IsArray(shape) && TrueRank(shape) == 0;
}
static bool IsScalarF32(const Shape& shape);
// Returns whether "shape" is a scalar (array) with the given element_type.
static bool IsScalarWithElementType(const Shape& shape,
PrimitiveType element_type);
// Extracts the size of the shape's dimension at dimension number
// GetDimensionNumber(dimension_number).