Home
last modified time | relevance | path

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

12

/external/tensorflow/tensorflow/compiler/xla/service/gpu/
Dgpu_transfer_manager.cc59 const Shape& sub_shape = ShapeUtil::GetSubshape(literal_shape, leaf.first); in TransferLiteralToInfeed()
122 const Shape& shape = ShapeUtil::GetSubshape(literal.shape(), leaf.first); in TransferLiteralFromOutfeed()
137 const Shape& shape = ShapeUtil::GetSubshape(literal.shape(), leaf.first); in TransferLiteralFromOutfeed()
Dinfeed_thunk.cc50 ShapeUtil::GetSubshape(source_buffers.shape(), shape_index); in ExecuteOnStream()
Dhlo_to_ir_bindings.cc144 ShapeUtil::GetSubshape(hlo.shape(), shape_index), ir_value, b_); in GetTypedIrValue()
199 ShapeUtil::GetSubshape(hlo.shape(), shape_index)); in GetIrArray()
Doutfeed_thunk.cc79 ShapeUtil::GetSubshape(output_buffers->shape(), shape_index); in ExecuteOnStream()
/external/tensorflow/tensorflow/compiler/xla/
Dshape_util_test.cc437 TEST(ShapeUtilTest, GetSubshape) { in TEST() argument
441 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(array_shape, {}))); in TEST()
449 ShapeUtil::Equal(tuple_shape, ShapeUtil::GetSubshape(tuple_shape, {}))); in TEST()
451 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {0}))); in TEST()
453 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {1}))); in TEST()
455 ShapeUtil::Equal(array_shape, ShapeUtil::GetSubshape(tuple_shape, {2}))); in TEST()
464 ShapeUtil::GetSubshape(nested_tuple_shape, {}))); in TEST()
466 array_shape, ShapeUtil::GetSubshape(nested_tuple_shape, {0}))); in TEST()
469 ShapeUtil::GetSubshape(nested_tuple_shape, {1}))); in TEST()
472 ShapeUtil::GetSubshape(nested_tuple_shape, {2, 0}))); in TEST()
[all …]
Dshape_tree.h689 ShapeUtil::GetSubshape(shape(), target_base_index), in CopySubtreeFrom()
690 ShapeUtil::GetSubshape(other.shape(), source_base_index))) in CopySubtreeFrom()
691 << ShapeUtil::GetSubshape(shape(), target_base_index) << " vs " in CopySubtreeFrom()
692 << ShapeUtil::GetSubshape(other.shape(), source_base_index); in CopySubtreeFrom()
Dliteral.cc409 elements.push_back(Literal(ShapeUtil::GetSubshape(shape(), {i}), in DecomposeTuple()
550 ShapeUtil::GetSubshape(shape(), dest_shape_index); in CopyFrom()
552 ShapeUtil::GetSubshape(src_literal.shape(), src_shape_index); in CopyFrom()
601 ShapeUtil::GetSubshape(shape(), dest_shape_index); in MoveFrom()
974 const Shape& subshape = ShapeUtil::GetSubshape(shape(), shape_index); in GetAsString()
1171 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in TupleToStringHelper()
1189 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in DenseArrayToStringHelper()
1276 const Shape& subshape = ShapeUtil::GetSubshape(literal.shape(), shape_index); in ToStringHelper()
1544 element.ConvertToShape(ShapeUtil::GetSubshape(dest_shape, {i}))); in ConvertToShape()
/external/tensorflow/tensorflow/compiler/xla/service/
Dbfloat16_normalization_test.cc274 EXPECT_EQ(ShapeUtil::GetSubshape(crs->shape(), {1}).element_type(), F32); in TEST_F()
298 EXPECT_EQ(ShapeUtil::GetSubshape(a2a->shape(), {0}).element_type(), BF16); in TEST_F()
299 EXPECT_EQ(ShapeUtil::GetSubshape(a2a->shape(), {1}).element_type(), BF16); in TEST_F()
327 EXPECT_EQ(ShapeUtil::GetSubshape(a2a->shape(), {0}).element_type(), F32); in TEST_F()
328 EXPECT_EQ(ShapeUtil::GetSubshape(a2a->shape(), {1}).element_type(), F32); in TEST_F()
361 EXPECT_EQ(ShapeUtil::GetSubshape(sort->shape(), {0}).element_type(), F32); in TEST_F()
386 EXPECT_EQ(ShapeUtil::GetSubshape(sort->shape(), {0}).element_type(), F32); in TEST_F()
Dlogical_buffer.h51 return ShapeUtil::GetSubshape(instruction_->shape(), index_); in shape()
Dbuffer_value.cc31 const Shape& shape = ShapeUtil::GetSubshape(instruction->shape(), index); in BufferValue()
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()
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.cc200 ShapeUtil::GetSubshape(param_shape, alias.parameter_index); in Verify()
202 ShapeUtil::GetSubshape(output_shape, output_index); in Verify()
Dlayout_assignment.cc776 ShapeUtil::GetSubshape(parameter->shape(), shape_index))) { in CheckParameterLayout()
813 ShapeUtil::GetSubshape(shape_with_layout, {i}); in CreateCopyWithNewLayout()
815 ShapeUtil::GetSubshape(instruction->shape(), {i}); in CreateCopyWithNewLayout()
950 ShapeUtil::GetSubshape(instruction->shape(), index); in CheckLayouts()
1268 ShapeUtil::GetSubshape(shape_layout.shape(), index).layout(), in PropagateUseConstraintToDefs()
1632 const Shape& subshape = ShapeUtil::GetSubshape(instruction->shape(), index); in InferArrayLayout()
1900 ? ShapeUtil::GetSubshape(instruction->literal().shape(), in RunOnComputation()
1956 const Shape& result_shape = ShapeUtil::GetSubshape( in RunOnComputation()
1986 ShapeUtil::GetSubshape(instruction->shape(), {0}).layout()); in ConstrainChannelLayouts()
2064 const auto& computed_subshape = ShapeUtil::GetSubshape( in PropagateComputationLayouts()
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()
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()
1112 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
1114 ShapeUtil::GetSubshape(while_hlo2->shape(), {0}), while_hlo2, 0)); in TEST_F()
1179 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
1181 ShapeUtil::GetSubshape(while_hlo2->shape(), {0}), while_hlo2, 0)); in TEST_F()
1479 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo1, 0)); in TEST_F()
1481 ShapeUtil::GetSubshape(while_hlo1->shape(), {0}), while_hlo2, 0)); in TEST_F()
Ddynamic_parameter_binding.cc126 ShapeUtil::GetSubshape( in Verify()
Dshaped_buffer.cc177 xla::ShapeUtil::GetSubshape(on_device_shape(), {index}); in TakeSubTree()
Dtransfer_manager.cc214 ShapeUtil::GetSubshape(*device_shape, index); in ReadDynamicShapes()
407 ShapeUtil::GetSubshape(shaped_buffer.on_device_shape(), index); in AllocateScopedShapedBuffer()
Dhlo_sharding.cc391 ShapeUtil::GetSubshape(shape, index_to_sharding.first), num_devices); in ValidateTuple()
633 ShapeUtil::GetLeafCount(ShapeUtil::GetSubshape(*sub_shape, {i})); in GetSubSharding()
635 sub_shape = &ShapeUtil::GetSubshape(*sub_shape, {idx}); in GetSubSharding()
/external/tensorflow/tensorflow/stream_executor/tpu/
Dtpu_executable_interface.cc101 int64 allocation_bytes = ShapeSize(ShapeUtil::GetSubshape( in AllocateOutputMemoryWithInputReuse()
146 ShapeUtil::GetSubshape(on_device_shape, result_index); in AllocateOutputMemoryWithInputReuse()
/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()
623 xla::ShapeUtil::GetSubshape(source.on_device_shape(), source_index); in AliasBufferFrom()
625 xla::ShapeUtil::GetSubshape(on_device_shape(), dest_index); in AliasBufferFrom()
/external/tensorflow/tensorflow/compiler/jit/
Dxla_tensor.cc55 xla::ShapeUtil::GetSubshape(on_device_shape, index_to_buffer.first); in AllocateShapedBuffer()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_execute_op.cc322 xla::ShapeUtil::GetSubshape(input_host_shape, {i}); in BuildComputationInputs()
422 xla::ShapeUtil::GetSubshape(scoped_buffers.on_host_shape(), {i}); in AllocateOutputTensors()
449 xla::ShapeUtil::GetSubshape(output_device_shape, {i}); in AllocateOutputTensors()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dutils.cc177 ShapeUtil::GetSubshape(sharded_shape, idx); in DetermineArgumentLayoutsFromCompileOptions()

12