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.cc435 TEST(ShapeUtilTest, GetSubshape) { in TEST() argument
439 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(array_shape, {}))); in TEST()
447 ShapeUtil::Equal(tuple_shape, ShapeUtil::GetSubshape(tuple_shape, {}))); in TEST()
449 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {0}))); in TEST()
451 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {1}))); in TEST()
453 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {2}))); in TEST()
462 ShapeUtil::GetSubshape(nested_tuple_shape, {}))); in TEST()
464 array_shape, ShapeUtil::GetSubshape(nested_tuple_shape, {0}))); in TEST()
467 ShapeUtil::GetSubshape(nested_tuple_shape, {1}))); in TEST()
470 ShapeUtil::GetSubshape(nested_tuple_shape, {2, 0}))); in TEST()
[all …]
Dliteral.cc341 elements.push_back(Literal(ShapeUtil::GetSubshape(shape(), {i}), in DecomposeTuple()
427 ShapeUtil::GetSubshape(shape(), dest_shape_index); in CopyFrom()
429 ShapeUtil::GetSubshape(src_literal.shape(), src_shape_index); in CopyFrom()
467 ShapeUtil::GetSubshape(shape(), dest_shape_index); in MoveFrom()
784 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); in GetAsString()
975 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in TupleToStringHelper()
993 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in DenseArrayToStringHelper()
1066 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in ToStringHelper()
1326 element.ConvertToShape(ShapeUtil::GetSubshape(dest_shape, {i}))); in ConvertToShape()
Dshape_tree.h650 ShapeUtil::GetSubshape(shape(), target_base_index), in CopySubtreeFrom()
651 ShapeUtil::GetSubshape(other.shape(), source_base_index))); in CopySubtreeFrom()
/external/tensorflow/tensorflow/compiler/xla/service/
Dbfloat16_normalization_test.cc273 EXPECT_EQ(ShapeUtil::GetSubshape(crs->shape(), {1}).element_type(), F32); in TEST_F()
297 EXPECT_EQ(ShapeUtil::GetSubshape(a2a->shape(), {0}).element_type(), BF16); in TEST_F()
298 EXPECT_EQ(ShapeUtil::GetSubshape(a2a->shape(), {1}).element_type(), BF16); in TEST_F()
326 EXPECT_EQ(ShapeUtil::GetSubshape(a2a->shape(), {0}).element_type(), F32); in TEST_F()
327 EXPECT_EQ(ShapeUtil::GetSubshape(a2a->shape(), {1}).element_type(), F32); in TEST_F()
360 EXPECT_EQ(ShapeUtil::GetSubshape(sort->shape(), {0}).element_type(), F32); in TEST_F()
385 EXPECT_EQ(ShapeUtil::GetSubshape(sort->shape(), {0}).element_type(), F32); in TEST_F()
Dlogical_buffer.h51 return ShapeUtil::GetSubshape(instruction_->shape(), index_); in shape()
Dshaped_buffer.cc175 xla::ShapeUtil::GetSubshape(on_host_shape(), {index}); in TakeSubTree()
177 xla::ShapeUtil::GetSubshape(on_device_shape(), {index}); in TakeSubTree()
Dbfloat16_normalization.cc136 ShapeUtil::GetSubshape(hlo->shape(), leaf_index); in ConvertType()
192 ShapeUtil::GetSubshape(original_shape, leaf_index); in ChangeOutputTypeThenInsertConvertBack()
271 CHECK(ShapeUtil::GetSubshape(hlo->shape(), {i}).IsArray()); in HandleMultipleOutputs()
273 output_types[i] = ShapeUtil::GetSubshape(hlo->shape(), {i}).element_type(); in HandleMultipleOutputs()
Dbuffer_value.cc31 const Shape& shape = ShapeUtil::GetSubshape(instruction->shape(), index); in BufferValue()
Dpattern_matcher_test.cc121 ShapeUtil::Equal(*subshape, ShapeUtil::GetSubshape(tuple_shape, {0}))); in TEST_F()
125 &ShapeUtil::GetSubshape(tuple_shape, {0}))))); in TEST_F()
129 &ShapeUtil::GetSubshape(tuple_shape, {1}))))); in TEST_F()
137 ShapeUtil::Equal(*subshape, ShapeUtil::GetSubshape(tuple_shape, {1}))); in TEST_F()
141 &ShapeUtil::GetSubshape(tuple_shape, {1}))))); in TEST_F()
145 &ShapeUtil::GetSubshape(tuple_shape, {0}))))); in TEST_F()
Dhlo_input_output_alias_config.cc201 ShapeUtil::GetSubshape(param_shape, alias.parameter_index); in Verify()
203 ShapeUtil::GetSubshape(output_shape, output_index); in Verify()
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()
745 ShapeUtil::GetSubshape(xla_while->shape(), {1}), xla_while, 1)); in BuildWhileInstruction_InitPointsToInterfering()
750 ShapeUtil::GetSubshape(xla_while->shape(), {0}), xla_while, 0)); in BuildWhileInstruction_InitPointsToInterfering()
913 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
915 ShapeUtil::GetSubshape(while_hlo2->shape(), {0}), while_hlo2, 0)); in TEST_F()
980 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
982 ShapeUtil::GetSubshape(while_hlo2->shape(), {0}), while_hlo2, 0)); in TEST_F()
1280 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
1282 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo2, 0)); in TEST_F()
Ddynamic_dimension_inference.h178 Shape subshape = ShapeUtil::GetSubshape(inst->shape(), index); in SetDynamicSize()
Dbfloat16_conversion_folding_test.cc266 EXPECT_EQ(ShapeUtil::GetSubshape(crs->shape(), {0}).element_type(), F32); in TEST_F()
267 EXPECT_EQ(ShapeUtil::GetSubshape(crs->shape(), {1}).element_type(), BF16); in TEST_F()
Dlayout_assignment.cc775 ShapeUtil::GetSubshape(parameter->shape(), shape_index))) { in CheckParameterLayout()
812 ShapeUtil::GetSubshape(shape_with_layout, {i}); in CreateCopyWithNewLayout()
814 ShapeUtil::GetSubshape(instruction->shape(), {i}); in CreateCopyWithNewLayout()
949 ShapeUtil::GetSubshape(instruction->shape(), index); in CheckLayouts()
1267 ShapeUtil::GetSubshape(shape_layout.shape(), index).layout(), in PropagateUseConstraintToDefs()
1609 const Shape& subshape = ShapeUtil::GetSubshape(instruction->shape(), index); in InferArrayLayout()
1877 ? ShapeUtil::GetSubshape(instruction->literal().shape(), in RunOnComputation()
1946 ShapeUtil::GetSubshape(instruction->shape(), {0}).layout()); in ConstrainChannelLayouts()
2024 const auto& computed_subshape = ShapeUtil::GetSubshape( in PropagateComputationLayouts()
Ddynamic_parameter_binding.cc126 ShapeUtil::GetSubshape( in Verify()
Dhlo_value.cc45 return ShapeUtil::GetSubshape(instruction->shape(), index); in shape()
Dbfloat16_propagation.cc209 const Shape& subshape = ShapeUtil::GetSubshape(hlo.shape(), index); in AllUsersConsumeBF16()
664 ShapeUtil::GetSubshape(hlo->shape(), leaf_index); in ResolveInconsistentFusions()
/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dcustom_call_thunk.cc95 ShapeUtil::GetSubshape(slices.shape(), index.ConsumeBack()) in AppendBuffersFor()
148 ShapeUtil::GetSubshape(result_slices_.shape(), index); in ExecuteOnStream()
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.cc166 ShapeUtil::GetSubshape(hlo.shape(), shape_index), module_); in GetTypedIrValue()
231 ShapeUtil::GetSubshape(hlo.shape(), shape_index)); in GetIrArray()
Dinfeed_thunk.cc62 const Shape& shape = ShapeUtil::GetSubshape(infeed_buffers.shape(), in ExecuteOnStream()
/external/tensorflow/tensorflow/compiler/xrt/
Dxrt_state.cc99 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first); in AllocateScopedShapedBuffer()
486 xla::ShapeUtil::GetSubshape(device_shape, index); in MakeTuple()
622 xla::ShapeUtil::GetSubshape(source.on_device_shape(), source_index); in AliasBufferFrom()
624 xla::ShapeUtil::GetSubshape(on_device_shape(), dest_index); in AliasBufferFrom()
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.cc59 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first); in AllocateShapedBuffer()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_transfer_manager.cc122 const Shape& tuple_element_shape = ShapeUtil::GetSubshape(shape, {i}); in TransferLiteralToInfeed()

12