Home
last modified time | relevance | path

Searched refs:GetTupleElementShape (Results 1 – 25 of 40) sorted by relevance

12

/external/tensorflow/tensorflow/compiler/xla/service/
Dgather_expander_test.cc106 ShapeUtil::GetTupleElementShape(while_shape, 1))); in TEST_F()
110 ShapeUtil::GetTupleElementShape(while_shape, 2))); in TEST_F()
114 ShapeUtil::GetTupleElementShape(while_shape, 3))); in TEST_F()
Dhlo_element_type_converter.cc72 Shape subshape = ShapeUtil::GetTupleElementShape(shape, i); in GetConvertedTupleShape()
93 const Shape& ele_shape = ShapeUtil::GetTupleElementShape(shape, i); in ConvertTupleElements()
96 const Shape& to_ele_shape = ShapeUtil::GetTupleElementShape(to_shape, i); in ConvertTupleElements()
Dhlo_domain_test.cc485 ShapeUtil::GetTupleElementShape(infeed_data->shape(), 0), infeed_data, in TEST_F()
492 ShapeUtil::GetTupleElementShape(infeed_data->shape(), 1), infeed_data, in TEST_F()
739 ShapeUtil::GetTupleElementShape(copy1->shape(), 0), copy1, 0)); in TEST_F()
742 ShapeUtil::GetTupleElementShape(tuple0->shape(), 1), tuple0, 1)); in TEST_F()
Dreduce_decomposer.cc90 output_shapes.push_back(ShapeUtil::GetTupleElementShape(shape, i)); in HandleReduce()
Dallocation_tracker.cc152 ShapeUtil::GetTupleElementShape(shaped_buffer->on_device_shape(), i), in DeconstructTuple()
Dtuple_simplifier_test.cc155 ShapeUtil::GetTupleElementShape(element->shape(), 0), element, 0)); in TEST_F()
Dhlo_verifier.cc1511 const Shape& dest_shape = ShapeUtil::GetTupleElementShape(operand_shape, 0); in HandleCopyDone()
1512 const Shape& src_shape = ShapeUtil::GetTupleElementShape(operand_shape, 1); in HandleCopyDone()
1522 return CheckShape(copy_done, ShapeUtil::GetTupleElementShape( in HandleCopyDone()
1542 {ShapeUtil::GetTupleElementShape(recv->shape(), 0), in HandleRecv()
1551 {ShapeUtil::GetTupleElementShape(recv_done->operand(0)->shape(), 0), in HandleRecvDone()
Dhlo_module_group_metadata.cc466 ShapeUtil::GetTupleElementShape(channel.recv_done->shape(), 0); in VerifyChannelInstructions()
Dbfloat16_propagation_test.cc953 EXPECT_EQ(ShapeUtil::GetTupleElementShape(domain->shape(), 0).element_type(), in TEST_F()
955 EXPECT_EQ(ShapeUtil::GetTupleElementShape(domain->shape(), 1).element_type(), in TEST_F()
Dlayout_assignment_test.cc82 const Layout& actual = ShapeUtil::GetTupleElementShape(shape, i).layout(); in ExpectTupleLayoutIs()
247 ShapeUtil::GetTupleElementShape(tuple->shape(), 1), constant1->shape())); in TEST_F()
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dtensor_list_utils.cc146 *buffer_shape = xla::ShapeUtil::GetTupleElementShape(list_shape, 0); in GetTensorListBufferShape()
213 auto buffer_shape = xla::ShapeUtil::GetTupleElementShape(list_shape, 0); in GetLeadingDimForTensorList()
233 xla::ShapeUtil::GetTupleElementShape(element_tensor_list_shape, i); in GetTensorListShapeFromElementTensorListShape()
280 xla::ShapeUtil::GetTupleElementShape(list_shape, i); in CreateZerosTensorListWithShape()
291 TF_RET_CHECK(xla::ShapeUtil::GetTupleElementShape(list_shape, tuple_size - 1) in CreateZerosTensorListWithShape()
375 xla::ShapeUtil::GetTupleElementShape(element_shape, i); in ExecuteTensorListPushBack()
437 xla::ShapeUtil::GetTupleElementShape(list_shape, i); in ExecuteTensorListPopBack()
518 xla::ShapeUtil::GetTupleElementShape(list_shape, 0); in ExecuteTensorListGetItem()
Dstack_ops.cc48 return XLAShapeToTensorShape(xla::ShapeUtil::GetTupleElementShape(shape, 0), in GetStackShape()
/external/tensorflow/tensorflow/compiler/xla/tests/
Dclient_test.cc102 ShapeUtil::GetTupleElementShape(result.shape(), 0), in XLA_TEST_F()
106 ShapeUtil::GetTupleElementShape(result.shape(), 1), in XLA_TEST_F()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_reshard_variables_op_util.cc115 xla::ShapeUtil::GetTupleElementShape(input_host_shape, i); in BuildInputBuffers()
278 xla::ShapeUtil::GetTupleElementShape(output_host_shape, i); in UpdateOutputVariables()
280 xla::ShapeUtil::GetTupleElementShape(output_device_shape, i); in UpdateOutputVariables()
Dtpu_compile_op_support.cc154 xla::ShapeUtil::GetTupleElementShape(tuple_shape_tree.shape(), in GetSubtree()
170 Shape element_shape = xla::ShapeUtil::GetTupleElementShape(shape, i); in GetPerDeviceShape()
222 xla::Shape shape = xla::ShapeUtil::GetTupleElementShape( in AddVariableUpdatesToCores()
285 xla::Shape shape = xla::ShapeUtil::GetTupleElementShape( in ComputeOutputShapesForEachCore()
Dtpu_execute_op.cc212 xla::ShapeUtil::GetTupleElementShape(input_host_shape, i); in BuildComputationInputs()
459 xla::ShapeUtil::GetTupleElementShape(output_device_shape, i); in AllocateOutputTensors()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dshape_util.cc34 xla::ShapeUtil::GetTupleElementShape(shape, i); in PopulateInfeedLayoutVector()
183 xla::ShapeUtil::GetTupleElementShape(input_shape, i); in GetShapeWithLayout()
Dxla_jit_compiled_cpu_function_test.cc228 const xla::Shape& result0 = ShapeUtil::GetTupleElementShape(result, 0); in TEST()
279 const xla::Shape& result0 = ShapeUtil::GetTupleElementShape(result, 0); in TEST()
/external/tensorflow/tensorflow/compiler/xla/service/cpu/
Dcpu_xfeed.cc254 ShapeUtil::GetTupleElementShape(literal.shape(), i); in TransferLiteralFromOutfeedOnCpu()
Dir_emitter.cc410 ShapeUtil::GetTupleElementShape(infeed->shape(), 0))); in HandleInfeed()
421 token_slice, ShapeUtil::GetTupleElementShape(infeed->shape(), 1)); in HandleInfeed()
437 ShapeUtil::GetTupleElementShape(data_shape, i); in HandleInfeed()
548 ShapeUtil::GetTupleElementShape(operand_shape, i); in HandleOutfeed()
2471 ShapeUtil::GetTupleElementShape(custom_call->shape(), i); in HandleCustomCall()
/external/tensorflow/tensorflow/compiler/aot/tests/
Dtfcompile_test.cc669 ShapeUtil::GetTupleElementShape(muladd_result, 0); in TEST()
672 ShapeUtil::GetTupleElementShape(muladd_result, 1); in TEST()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.h62 CreateConstantBase(ShapeUtil::GetTupleElementShape(shape, i), in CreateConstantBase()
Dspmd_partitioner_util.cc77 if (!EvenlyPartitions(ShapeUtil::GetTupleElementShape(shape, i), in EvenlyPartitions()
102 MakePartitionedShape(ShapeUtil::GetTupleElementShape(shape, i), in MakePartitionedShape()
124 ShapeUtil::GetTupleElementShape(shape, i), in MakeNonPaddedShapeForGivenPartition()
/external/tensorflow/tensorflow/compiler/xla/
Dshape_util.h434 static const Shape& GetTupleElementShape(const Shape& shape, int64_t index);
Dliteral_test.cc2026 ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 0).ToProto(); in TEST_F()
2044 ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 0).ToProto(); in TEST_F()
2049 ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 1).ToProto(); in TEST_F()

12