Home
last modified time | relevance | path

Searched refs:tuple_shapes (Results 1 – 25 of 99) sorted by relevance

1234

/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/
H A Dtpu_infeed_test.py32 self.assertEqual(i.tuple_shapes, None)
39 self.assertEqual(i.tuple_shapes, None)
41 i = tpu_feed.InfeedQueue(tuple_shapes=[[1], [2, 3]])
44 self.assertEqual(i.tuple_shapes, [[1], [2, 3]])
49 self.assertEqual(i.tuple_shapes, None)
58 i = tpu_feed.InfeedQueue(number_of_tuple_elements=2, tuple_shapes=[[1]])
62 i = tpu_feed.InfeedQueue(tuple_shapes=[[1], [2, 3]], shard_dimensions=[1])
74 self.assertEqual(i.tuple_shapes, [[1], [2, 3]])
76 self.assertEqual(i.tuple_shapes, [[1, 2], [3, 4]])
86 self.assertEqual(i.tuple_shapes, [[6], [3, 18]])
[all …]
H A Dtpu_feed.py131 tuple_shapes=None, argument
179 elif tuple_shapes is not None:
180 number_of_tuple_elements = len(tuple_shapes)
198 if tuple_shapes is not None:
199 self.set_tuple_shapes(tuple_shapes)
212 if self.tuple_shapes is not None:
264 def tuple_shapes(self): member in InfeedQueue
268 def set_tuple_shapes(self, tuple_shapes): argument
284 if len(tuple_shapes) != self.number_of_tuple_elements:
290 tuple_shapes = [tensor_shape.as_shape(shape) for shape in tuple_shapes]
[all …]
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/
H A Dlayout_util_test.cc246 EXPECT_TRUE(shape.tuple_shapes(0).has_layout()); in TEST_F()
247 EXPECT_TRUE(shape.tuple_shapes(2).tuple_shapes(1).has_layout()); in TEST_F()
252 EXPECT_FALSE(shape.tuple_shapes(0).has_layout()); in TEST_F()
253 EXPECT_FALSE(shape.tuple_shapes(2).tuple_shapes(1).has_layout()); in TEST_F()
273 EXPECT_FALSE(LayoutUtil::Equal(shape.tuple_shapes(0).layout(), in TEST_F()
274 shape.tuple_shapes(1).layout())); in TEST_F()
276 EXPECT_TRUE(LayoutUtil::Equal(shape.tuple_shapes(0).layout(), in TEST_F()
277 shape.tuple_shapes(1).layout())); in TEST_F()
279 LayoutUtil::GetDefaultLayoutForShape(shape.tuple_shapes(0)), in TEST_F()
280 shape.tuple_shapes(1).layout())); in TEST_F()
H A Dshape_util.cc329 for (const Shape& s : shape.tuple_shapes()) { in MoveDimToMajor()
549 absl::c_any_of(shape.tuple_shapes(), in IsNestedTuple()
566 TF_DCHECK_OK(ValidateShapeWithOptionalLayout(shape.tuple_shapes(index))); in GetTupleElementShape()
567 return shape.tuple_shapes(index); in GetTupleElementShape()
584 std::vector<Shape> new_elements(tuple.tuple_shapes().begin() + start, in SliceTuple()
585 tuple.tuple_shapes().begin() + limit); in SliceTuple()
614 for (const Shape& element_shape : shape.tuple_shapes()) { in ElementsInRecursive()
625 for (const Shape& element_shape : shape.tuple_shapes()) { in HasPrimitiveType()
645 const auto& tuple_shapes = shape.tuple_shapes(); in HumanString() local
646 for (int64_t i = 0; i < tuple_shapes.size(); ++i) { in HumanString()
[all …]
H A Dshape.cc50 for (const ShapeProto& element_shape : shape_proto.tuple_shapes()) { in Shape()
75 for (const Shape& shape : tuple_shapes()) { in ToProto()
147 lhs.tuple_shapes(), rhs.tuple_shapes(), in operator ()()
H A Dshape.h44 std::vector<Shape> tuple_shapes) in Shape() argument
49 tuple_shapes_(std::move(tuple_shapes)) {} in Shape()
145 const Shape& tuple_shapes(int index) const { return tuple_shapes_.at(index); } in tuple_shapes() function
152 const std::vector<Shape>& tuple_shapes() const { return tuple_shapes_; } in tuple_shapes() function
H A Dlayout_util.cc182 for (auto& element_shape : shape.tuple_shapes()) { in ValidateLayoutInShape()
362 return absl::c_all_of(shape.tuple_shapes(), in HasLayout()
440 TF_RETURN_IF_ERROR(CopyLayoutInternal(src.tuple_shapes(i), in CopyLayoutInternal()
473 if (!LayoutsInShapesEqual(lhs.tuple_shapes(i), rhs.tuple_shapes(i))) { in LayoutsInShapesEqual()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/client/lib/
H A Dtesting.cc41 for (const Shape& s : shape.tuple_shapes()) { in DataSizeOfShape()
55 const auto& tuple_shapes = shape.tuple_shapes(); in BuildFakeDataOpOnDevice() local
56 parts.reserve(tuple_shapes.size()); in BuildFakeDataOpOnDevice()
57 for (const Shape& s : tuple_shapes) { in BuildFakeDataOpOnDevice()
H A Ddynamic_shaped_ops.cc49 subshapes.push_back(&shapes[j]->tuple_shapes(i)); in FindMaxShape()
84 left_branch_shape.tuple_shapes(i), in ReconsileBranchDifference()
85 right_branch_shape.tuple_shapes(i), sub_tuple)); in ReconsileBranchDifference()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/stream_executor/tpu/
H A Dc_api_conversions.cc241 c_shape->tuple_shapes = new XLA_Shape[c_shape->ntuple_shapes]; in ToC()
243 ToC(xla_shape.tuple_shapes(i), &c_shape->tuple_shapes[i]); in ToC()
259 std::vector<xla::Shape> tuple_shapes; in FromC() local
260 tuple_shapes.reserve(c_shape->ntuple_shapes); in FromC()
262 tuple_shapes.push_back(FromC(&c_shape->tuple_shapes[i])); in FromC()
266 dims, dynamic_dims, std::move(tuple_shapes)); in FromC()
282 Destroy(&c_shape->tuple_shapes[i]); in Destroy()
284 delete[] c_shape->tuple_shapes; in Destroy()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dwhile_loop_simplifier.cc81 &while_init->shape().tuple_shapes(old_idx)); in RemoveDeadTupleIndices()
165 while_init->shape().tuple_shapes(old_idx), while_init, old_idx))); in RemoveDeadTupleIndices()
204 new_while_op->shape().tuple_shapes(gte_idx), new_while_op, in RemoveDeadTupleIndices()
209 while_init->shape().tuple_shapes(old_idx), while_init, old_idx))); in RemoveDeadTupleIndices()
487 comp->parameter_instruction(0)->shape().tuple_shapes(tuple_index), in TryRemoveRepeatedWhileTupleIndicesHelper()
666 new_while_shape_elems.push_back(&while_shape.tuple_shapes(i)); in TryRemoveConstantParams()
690 while_shape.tuple_shapes(i), instr, i))); in TryRemoveConstantParams()
707 while_shape.tuple_shapes(i), instr, j))); in TryRemoveConstantParams()
948 const Shape& subshape = desired_shape.tuple_shapes(i); in UnflattenTupleInstr()
966 desired_shape.tuple_shapes(i), new_instrs); in UnflattenTupleInstr()
[all …]
H A Dhlo_verifier.cc354 : ag->shape().tuple_shapes(i); in CheckCommonAllGatherInvariants()
358 ? ag->shape().tuple_shapes(1) in CheckCommonAllGatherInvariants()
359 : ag->shape().tuple_shapes(1).tuple_shapes(i); in CheckCommonAllGatherInvariants()
439 : ars->shape().tuple_shapes(i); in HandleReduceScatter()
444 (ars->operand_count() == 1) ? ars->shape() : ars->shape().tuple_shapes(0); in HandleReduceScatter()
529 absl::c_all_of(buffer_offset_shape.tuple_shapes(), in CheckBufferOffset()
532 absl::c_all_of(buffer_offset_shape.tuple_shapes(), in CheckBufferOffset()
541 if (absl::c_any_of(buffer_offset_shape.tuple_shapes(), in CheckBufferOffset()
597 CheckBufferOffset(input_buffer_shape.tuple_shapes(i), in CheckInplaceCollectivePermute()
598 input_offset_shape.tuple_shapes(i)); in CheckInplaceCollectivePermute()
[all …]
H A Dtuple_util.cc34 input_shape.tuple_shapes(i), input_tuple, i))); in ExtractPrefix()
53 input_shape.tuple_shapes(i), input_tuple, i))); in AppendSuffix()
94 const Shape& subshape = tuple_shape.tuple_shapes(i); in ReplaceTupleWith()
167 instruction->shape().tuple_shapes(index), instruction, index)); in AddGetTupleElements()
H A Dconditional_simplifier.cc102 new_tuple_shapes.push_back(&param->shape().tuple_shapes(i)); in TryRemoveUnusedConditionalOperands()
138 old_shape.tuple_shapes(i), in TryRemoveUnusedConditionalOperands()
252 &old_shape.tuple_shapes(new_to_old_mapping[new_index])); in RemoveUnusedTupleElements()
293 old_shape.tuple_shapes(old_index), old_root, old_index))); in RemoveUnusedTupleElements()
544 hlo->shape().tuple_shapes(i), hlo, i)); in TryRemoveConditional()
H A Dconditional_simplifier_test.cc214 EXPECT_EQ(conditional->operand(1)->shape().tuple_shapes().size(), 2); in TEST_F()
215 EXPECT_EQ(conditional->operand(2)->shape().tuple_shapes().size(), 2); in TEST_F()
219 call->to_apply()->parameter_instruction(0)->shape().tuple_shapes().size(), in TEST_F()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/
H A Dcustom_call_handler.cc92 const int64_t batch_size = hlo->shape().tuple_shapes(0).dimensions(batch_dim); in HandleCustomCallTopK()
93 const int64_t k = hlo->shape().tuple_shapes(0).dimensions(sort_dim); in HandleCustomCallTopK()
144 topk->shape().tuple_shapes(0), topk, 0)); in HandleCustomCallTopK()
148 value_gte, partitioned_topk.base_shape().tuple_shapes(0), in HandleCustomCallTopK()
157 topk->shape().tuple_shapes(1), topk, 1)); in HandleCustomCallTopK()
175 index_gte, partitioned_topk.base_shape().tuple_shapes(1), in HandleCustomCallTopK()
217 replicated_sort.hlo()->shape().tuple_shapes(0), replicated_sort.hlo(), in HandleCustomCallTopK()
221 replicated_sort.hlo()->shape().tuple_shapes(1), replicated_sort.hlo(), in HandleCustomCallTopK()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/
H A Dcudnn_pad_for_convolutions.cc74 CHECK_EQ(0, conv->shape().tuple_shapes(1).dimensions(0)) in PadConv()
83 const Shape& result_shape = conv->shape().tuple_shapes(0); in PadConv()
187 const Shape& result_shape = conv->shape().tuple_shapes(0); in TryResolvePaddedShapesForTensorCore()
203 new_result_shape = conv->shape().tuple_shapes(0); in TryResolvePaddedShapesForTensorCore()
301 const Shape& result_shape = conv->shape().tuple_shapes(0); in TryResolvePaddedShapesForIntegerConvolution()
320 new_result_shape = conv->shape().tuple_shapes(0); in TryResolvePaddedShapesForIntegerConvolution()
H A Dgpu_hlo_cost_analysis.cc76 options_.shape_size(custom_call->shape().tuple_shapes(0))); in HandleCustomCall()
95 return shape.tuple_shapes(0); in GetConvolutionFlops()
H A Dtree_reduction_rewriter.cc244 absl::InlinedVector<Shape, 2> tuple_shapes; in RewriteBatchDimensionLargerThanTile() local
246 tuple_shapes.push_back( in RewriteBatchDimensionLargerThanTile()
252 ShapeUtil::MakeMaybeTupleShape(tuple_shapes), hlo->inputs(), in RewriteBatchDimensionLargerThanTile()
H A Dreduction_degenerate_dim_remover.cc50 ? instr->shape().tuple_shapes(idx) in HandleReduce()
107 orig_reduce_shape.tuple_shapes(oidx), gte))); in HandleReduce()
H A Dcudnn_fused_conv_rewriter.cc121 const Shape& shape = conv->shape().tuple_shapes(0); in EnsureIsConvBiasActivation()
375 PrimitiveType conv_ty = conv->shape().tuple_shapes(0).element_type(); in FuseSideInputAlpha()
635 if (is_integral_not_s8(instr->shape().tuple_shapes(0)) || in CheckNoIllegalIntegerConvs()
703 PrimitiveType conv_out_ty = instr->shape().tuple_shapes(0).element_type(); in VlogStats()
H A Dreduction_layout_normalizer.cc73 reduce->shape().IsTuple() ? reduce->shape().tuple_shapes(operand_idx) in HandleReduce()
171 orig_reduce_shape.tuple_shapes(oidx), gte))); in HandleReduce()
H A Dgpu_conv_algorithm_picker.cc445 const Shape& result_shape = instr->shape().tuple_shapes(0); in AutotuneOneConvRunner()
649 const Shape& result_shape = instr->shape().tuple_shapes(0); in PickBestAlgorithmNoCacheCuda()
825 ShapeUtil::ByteSizeOf(instr->shape().tuple_shapes(0)))); in PickBestAlgorithmNoCacheRocm()
947 {instr->shape().tuple_shapes(0), in RunOnInstruction()
974 new_call_shape.tuple_shapes(0), new_call, 0)), in RunOnInstruction()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/kernels/
H A Dwhile_op.cc185 body.xla_output_shape.tuple_shapes(i); in VerifyBodyInputAndOutputShapeMatch()
191 body.xla_output_shape.tuple_shapes(ctx->num_inputs()); in VerifyBodyInputAndOutputShapeMatch()
428 arg.shape = body_output_shape.tuple_shapes(i); in Compile()
520 const xla::Shape& list_shape = body_input_shape.tuple_shapes(i); in Compile()
531 const xla::Shape& shape = list_shape.tuple_shapes(i); in Compile()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python_api/
H A Dxla_shape.py56 component_message = self.message.tuple_shapes.add()
75 def tuple_shapes(self): member in Shape

1234