Home
last modified time | relevance | path

Searched refs:GetTupleElementShape (Results 1 – 25 of 36) 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.cc70 Shape subshape = ShapeUtil::GetTupleElementShape(shape, i); in GetConvertedTupleShape()
91 const Shape& ele_shape = ShapeUtil::GetTupleElementShape(shape, i); in ConvertTupleElements()
94 const Shape& to_ele_shape = ShapeUtil::GetTupleElementShape(to_shape, i); in ConvertTupleElements()
Dhlo_domain_test.cc392 ShapeUtil::GetTupleElementShape(infeed_data->shape(), 0), infeed_data, in TEST_F()
399 ShapeUtil::GetTupleElementShape(infeed_data->shape(), 1), infeed_data, in TEST_F()
646 ShapeUtil::GetTupleElementShape(copy1->shape(), 0), copy1, 0)); in TEST_F()
649 ShapeUtil::GetTupleElementShape(tuple0->shape(), 1), tuple0, 1)); in TEST_F()
Dallocation_tracker.cc151 ShapeUtil::GetTupleElementShape(shaped_buffer->on_device_shape(), i), in DeconstructTuple()
Dhlo_verifier.cc385 hlo, ShapeUtil::GetTupleElementShape(hlo->operand(0)->shape(), 0)); in HandleCollectivePermuteDone()
993 const Shape& dest_shape = ShapeUtil::GetTupleElementShape(operand_shape, 0); in HandleCopyDone()
994 const Shape& src_shape = ShapeUtil::GetTupleElementShape(operand_shape, 1); in HandleCopyDone()
1004 return CheckShape(copy_done, ShapeUtil::GetTupleElementShape( in HandleCopyDone()
1024 {ShapeUtil::GetTupleElementShape(recv->shape(), 0), in HandleRecv()
1033 {ShapeUtil::GetTupleElementShape(recv_done->operand(0)->shape(), 0), in HandleRecvDone()
Dtuple_simplifier_test.cc157 ShapeUtil::GetTupleElementShape(element->shape(), 0), element, 0)); in TEST_F()
Dhlo_module_group_metadata.cc466 ShapeUtil::GetTupleElementShape(channel.recv_done->shape(), 0); in VerifyChannelInstructions()
Dlayout_assignment_test.cc81 const Layout& actual = ShapeUtil::GetTupleElementShape(shape, i).layout(); in ExpectTupleLayoutIs()
246 ShapeUtil::GetTupleElementShape(tuple->shape(), 1), constant1->shape())); in TEST_F()
Dbfloat16_propagation_test.cc996 EXPECT_EQ(ShapeUtil::GetTupleElementShape(domain->shape(), 0).element_type(), in TEST_F()
998 EXPECT_EQ(ShapeUtil::GetTupleElementShape(domain->shape(), 1).element_type(), in TEST_F()
/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
Dtensor_list_utils.cc146 *buffer_shape = xla::ShapeUtil::GetTupleElementShape(list_shape, 0); in GetTensorListBufferShape()
212 auto buffer_shape = xla::ShapeUtil::GetTupleElementShape(list_shape, 0); in GetLeadingDimForTensorList()
232 xla::ShapeUtil::GetTupleElementShape(element_tensor_list_shape, i); in GetTensorListShapeFromElementTensorListShape()
278 xla::ShapeUtil::GetTupleElementShape(list_shape, i); in CreateZerosTensorListWithShape()
289 TF_RET_CHECK(xla::ShapeUtil::GetTupleElementShape(list_shape, tuple_size - 1) in CreateZerosTensorListWithShape()
373 xla::ShapeUtil::GetTupleElementShape(element_shape, i); in ExecuteTensorListPushBack()
435 xla::ShapeUtil::GetTupleElementShape(list_shape, i); in ExecuteTensorListPopBack()
516 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.cc103 ShapeUtil::GetTupleElementShape(result.shape(), 0), in XLA_TEST_F()
107 ShapeUtil::GetTupleElementShape(result.shape(), 1), in XLA_TEST_F()
/external/tensorflow/tensorflow/core/tpu/kernels/
Dtpu_compile_op_support.cc152 xla::ShapeUtil::GetTupleElementShape(tuple_shape_tree.shape(), in GetSubtree()
168 Shape element_shape = xla::ShapeUtil::GetTupleElementShape(shape, i); in GetPerDeviceShape()
220 xla::Shape shape = xla::ShapeUtil::GetTupleElementShape( in AddVariableUpdatesToCores()
280 xla::Shape shape = xla::ShapeUtil::GetTupleElementShape( in ComputeOutputShapesForEachCore()
Dtpu_execute_op.cc213 xla::ShapeUtil::GetTupleElementShape(input_host_shape, i); in BuildComputationInputs()
460 xla::ShapeUtil::GetTupleElementShape(output_device_shape, i); in AllocateOutputTensors()
/external/tensorflow/tensorflow/compiler/tf2xla/
Dshape_util.cc34 xla::ShapeUtil::GetTupleElementShape(shape, i); in PopulateInfeedLayoutVector()
184 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_transfer_manager.cc208 ShapeUtil::GetTupleElementShape(literal.shape(), i); in TransferLiteralFromOutfeed()
Dir_emitter.cc397 ShapeUtil::GetTupleElementShape(infeed->shape(), 0))); in HandleInfeed()
407 token_slice, ShapeUtil::GetTupleElementShape(infeed->shape(), 1)); in HandleInfeed()
423 ShapeUtil::GetTupleElementShape(data_shape, i); in HandleInfeed()
534 ShapeUtil::GetTupleElementShape(operand_shape, i); in HandleOutfeed()
2339 ShapeUtil::GetTupleElementShape(custom_call->shape(), i); in HandleCustomCall()
/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner_util.cc60 ShapeUtil::GetTupleElementShape(shape, i), value.Clone(), b)); in CreateConstant()
76 CreateZero(ShapeUtil::GetTupleElementShape(shape, i), b)); in CreateZero()
97 CreateOne(ShapeUtil::GetTupleElementShape(shape, i), b)); in CreateOne()
130 if (!EvenlyPartitions(ShapeUtil::GetTupleElementShape(shape, i), in EvenlyPartitions()
153 MakePartitionedShape(ShapeUtil::GetTupleElementShape(shape, i), in MakePartitionedShape()
173 ShapeUtil::GetTupleElementShape(shape, i), in MakeNonPaddedShapeForGivenPartition()
Dspmd_partitioner.cc389 auto subshape = ShapeUtil::GetTupleElementShape(shape, i); in ReshardNoCache()
395 element, ShapeUtil::GetTupleElementShape(base_shape_, i), state_) in ReshardNoCache()
2478 ShapeUtil::GetTupleElementShape(tuple.hlo()->shape(), hlo->tuple_index()), in HandleGetTupleElement()
2492 const Shape& shape = ShapeUtil::GetTupleElementShape(hlo->shape(), 0); in HandleInfeed()
2777 ShapeUtil::GetTupleElementShape(reduce_shape, i), local_reduce, in HandleReduce()
/external/tensorflow/tensorflow/compiler/aot/tests/
Dtfcompile_test.cc652 ShapeUtil::GetTupleElementShape(muladd_result, 0); in TEST()
655 ShapeUtil::GetTupleElementShape(muladd_result, 1); in TEST()
/external/tensorflow/tensorflow/compiler/xla/
Dshape_util.h535 static const Shape& GetTupleElementShape(const Shape& shape, int64 index);
Dliteral_test.cc2017 ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 0).ToProto(); in TEST_F()
2035 ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 0).ToProto(); in TEST_F()
2040 ShapeUtil::GetTupleElementShape(Shape(proto.shape()), 1).ToProto(); in TEST_F()
Dshape_util.cc500 /* static */ const Shape& ShapeUtil::GetTupleElementShape(const Shape& shape, in GetTupleElementShape() function in xla::ShapeUtil
993 ShapeUtil::GetTupleElementShape(shape, i), func, index)); in ForEachSubshapeHelper()
Dliteral_comparison.cc208 total += RecursiveElementCount(ShapeUtil::GetTupleElementShape(shape, i)); in RecursiveElementCount()

12