/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/ |
H A D | tpu_infeed_test.py | 32 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 D | tpu_feed.py | 131 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 D | layout_util_test.cc | 246 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 D | shape_util.cc | 329 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 D | shape.cc | 50 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 D | shape.h | 44 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 D | layout_util.cc | 182 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 D | testing.cc | 41 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 D | dynamic_shaped_ops.cc | 49 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 D | c_api_conversions.cc | 241 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 D | while_loop_simplifier.cc | 81 &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 D | hlo_verifier.cc | 354 : 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 D | tuple_util.cc | 34 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 D | conditional_simplifier.cc | 102 new_tuple_shapes.push_back(¶m->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 D | conditional_simplifier_test.cc | 214 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 D | custom_call_handler.cc | 92 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 D | cudnn_pad_for_convolutions.cc | 74 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 D | gpu_hlo_cost_analysis.cc | 76 options_.shape_size(custom_call->shape().tuple_shapes(0))); in HandleCustomCall() 95 return shape.tuple_shapes(0); in GetConvolutionFlops()
|
H A D | tree_reduction_rewriter.cc | 244 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 D | reduction_degenerate_dim_remover.cc | 50 ? instr->shape().tuple_shapes(idx) in HandleReduce() 107 orig_reduce_shape.tuple_shapes(oidx), gte))); in HandleReduce()
|
H A D | cudnn_fused_conv_rewriter.cc | 121 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 D | reduction_layout_normalizer.cc | 73 reduce->shape().IsTuple() ? reduce->shape().tuple_shapes(operand_idx) in HandleReduce() 171 orig_reduce_shape.tuple_shapes(oidx), gte))); in HandleReduce()
|
H A D | gpu_conv_algorithm_picker.cc | 445 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 D | while_op.cc | 185 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 D | xla_shape.py | 56 component_message = self.message.tuple_shapes.add() 75 def tuple_shapes(self): member in Shape
|