Home
last modified time | relevance | path

Searched refs:GetSubshape (Results 1 – 25 of 39) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/
Dshape_util_test.cc413 TEST(ShapeUtilTest, GetSubshape) { in TEST() argument
417 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(array_shape, {}))); in TEST()
425 ShapeUtil::Equal(tuple_shape, ShapeUtil::GetSubshape(tuple_shape, {}))); in TEST()
427 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {0}))); in TEST()
429 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {1}))); in TEST()
431 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {2}))); in TEST()
440 ShapeUtil::GetSubshape(nested_tuple_shape, {}))); in TEST()
442 array_shape, ShapeUtil::GetSubshape(nested_tuple_shape, {0}))); in TEST()
445 ShapeUtil::GetSubshape(nested_tuple_shape, {1}))); in TEST()
448 ShapeUtil::GetSubshape(nested_tuple_shape, {2, 0}))); in TEST()
[all …]
Dliteral.cc342 elements.push_back(Literal(ShapeUtil::GetSubshape(shape(), {i}), in DecomposeTuple()
431 ShapeUtil::GetSubshape(shape(), dest_shape_index); in CopyFrom()
433 ShapeUtil::GetSubshape(src_literal.shape(), src_shape_index); in CopyFrom()
471 ShapeUtil::GetSubshape(shape(), dest_shape_index); in MoveFrom()
790 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); in GetAsString()
835 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); in GetSparseElementAsString()
1054 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in TupleToStringHelper()
1098 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in DenseArrayToStringHelper()
1171 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in ToStringHelper()
1439 element.ConvertToShape(ShapeUtil::GetSubshape(dest_shape, {i}))); in ConvertToShape()
Dshape_tree.h647 ShapeUtil::GetSubshape(shape(), target_base_index), in CopySubtreeFrom()
648 ShapeUtil::GetSubshape(other.shape(), source_base_index))); in CopySubtreeFrom()
/external/tensorflow/tensorflow/compiler/xla/service/
Dshaped_buffer.cc162 xla::ShapeUtil::GetSubshape(on_host_shape(), {index}); in TakeSubTree()
164 xla::ShapeUtil::GetSubshape(on_device_shape(), {index}); in TakeSubTree()
Dbuffer_liveness_test.cc72 ShapeUtil::Compatible(ShapeUtil::GetSubshape(a->shape(), index), in TupleElementsMayInterfere()
73 ShapeUtil::GetSubshape(b->shape(), index))); in TupleElementsMayInterfere()
492 ShapeUtil::GetSubshape(tuple_param0->shape(), {0}); in TEST_F()
504 ShapeUtil::GetSubshape(tuple_param0->shape(), {1}); in TEST_F()
556 ShapeUtil::GetSubshape(tuple_param0->shape(), {0}); in TEST_F()
568 ShapeUtil::GetSubshape(tuple_param0->shape(), {1}); in TEST_F()
Dlogical_buffer.h51 return ShapeUtil::GetSubshape(instruction_->shape(), index_); in shape()
Dpattern_matcher_test.cc145 ShapeUtil::Equal(*subshape, ShapeUtil::GetSubshape(tuple_shape, {0}))); in TEST()
149 &ShapeUtil::GetSubshape(tuple_shape, {0}))))); in TEST()
153 &ShapeUtil::GetSubshape(tuple_shape, {1}))))); in TEST()
161 ShapeUtil::Equal(*subshape, ShapeUtil::GetSubshape(tuple_shape, {1}))); in TEST()
165 &ShapeUtil::GetSubshape(tuple_shape, {1}))))); in TEST()
169 &ShapeUtil::GetSubshape(tuple_shape, {0}))))); in TEST()
Dbuffer_value.cc31 const Shape& shape = ShapeUtil::GetSubshape(instruction->shape(), index); in BufferValue()
Dhlo_input_output_alias_config.cc195 ShapeUtil::GetSubshape(param_shape, alias.parameter_index); in Verify()
197 ShapeUtil::GetSubshape(output_shape, output_index); in Verify()
Dbfloat16_normalization_test.cc271 EXPECT_EQ(ShapeUtil::GetSubshape(crs->shape(), {1}).element_type(), F32); in TEST_F()
301 EXPECT_EQ(ShapeUtil::GetSubshape(sort->shape(), {0}).element_type(), F32); in TEST_F()
326 EXPECT_EQ(ShapeUtil::GetSubshape(sort->shape(), {0}).element_type(), F32); in TEST_F()
Dhlo_sharding.cc283 ShapeUtil::GetSubshape(shape, index_to_sharding.first), num_devices); in ValidateTuple()
454 ShapeUtil::GetLeafCount(ShapeUtil::GetSubshape(*sub_shape, {i})); in GetSubSharding()
456 sub_shape = &ShapeUtil::GetSubshape(*sub_shape, {idx}); in GetSubSharding()
Dcopy_insertion_test.cc352 ShapeUtil::GetSubshape(param->shape(), {0}), param, 0)); in TEST_F()
389 ShapeUtil::GetSubshape(select->shape(), {0}), select, 0)); in TEST_F()
739 ShapeUtil::GetSubshape(xla_while->shape(), {1}), xla_while, 1)); in BuildWhileInstruction_InitPointsToInterfering()
744 ShapeUtil::GetSubshape(xla_while->shape(), {0}), xla_while, 0)); in BuildWhileInstruction_InitPointsToInterfering()
907 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
909 ShapeUtil::GetSubshape(while_hlo2->shape(), {0}), while_hlo2, 0)); in TEST_F()
969 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
971 ShapeUtil::GetSubshape(while_hlo2->shape(), {0}), while_hlo2, 0)); in TEST_F()
1269 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
1271 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo2, 0)); in TEST_F()
Dbfloat16_conversion_folding_test.cc265 EXPECT_EQ(ShapeUtil::GetSubshape(crs->shape(), {0}).element_type(), F32); in TEST_F()
266 EXPECT_EQ(ShapeUtil::GetSubshape(crs->shape(), {1}).element_type(), BF16); in TEST_F()
Ddynamic_parameter_binding.cc126 ShapeUtil::GetSubshape( in Verify()
Dlayout_assignment.cc786 ShapeUtil::GetSubshape(shape_with_layout, {i}); in CreateCopyWithNewLayout()
788 ShapeUtil::GetSubshape(instruction->shape(), {i}); in CreateCopyWithNewLayout()
897 ShapeUtil::GetSubshape(instruction->shape(), index); in CheckLayouts()
1210 ShapeUtil::GetSubshape(shape_layout.shape(), index).layout(), in PropagateUseConstraintToDefs()
1491 const Shape& subshape = ShapeUtil::GetSubshape(instruction->shape(), index); in InferArrayLayout()
1760 ? ShapeUtil::GetSubshape(instruction->literal().shape(), in RunOnComputation()
1824 ShapeUtil::GetSubshape(instruction->shape(), {0}).layout()); in ConstrainChannelLayouts()
Dhlo_value.cc44 return ShapeUtil::GetSubshape(instruction->shape(), index); in shape()
Dbfloat16_propagation.cc209 const Shape& subshape = ShapeUtil::GetSubshape(hlo.shape(), index); in AllUsersConsumeBF16()
648 ShapeUtil::GetSubshape(hlo->shape(), leaf_index); in ResolveInconsistentFusions()
Dtuple_points_to_analysis.cc712 ShapeUtil::GetSubshape(operand->shape(), operand_index); in CanShareOperandBufferWithUser()
714 ShapeUtil::GetSubshape(user->shape(), user_index); in CanShareOperandBufferWithUser()
Dtransfer_manager.cc328 const Shape& subshape = ShapeUtil::GetSubshape(on_device_shape, index); in AllocateScopedShapedBuffer()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_transfer_manager.cc127 const Shape& shape = ShapeUtil::GetSubshape(shape_tree->shape(), *index); in ShapeTreeToLiteral()
157 const Shape& shape = ShapeUtil::GetSubshape(literal_shape, index); in TransferLiteralFromOutfeed()
Dhlo_to_ir_bindings.cc157 ShapeUtil::GetSubshape(hlo.shape(), shape_index), module_); in GetTypedIrValue()
221 ShapeUtil::GetSubshape(hlo.shape(), shape_index)); in GetIrArray()
Dinfeed_thunk.cc60 const Shape& shape = ShapeUtil::GetSubshape(infeed_buffers.shape(), in ExecuteOnStream()
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.cc59 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first); in AllocateShapedBuffer()
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt_state.cc113 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first); in AllocateScopedShapedBuffer()
412 xla::ShapeUtil::GetSubshape(device_shape, index); in MakeTuple()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_transfer_manager.cc122 const Shape& tuple_element_shape = ShapeUtil::GetSubshape(shape, {i}); in TransferLiteralToInfeed()

12