/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/ |
H A D | tuple_util.cc | 29 std::vector<HloInstruction*> tuple_elements; in ExtractPrefix() local 30 tuple_elements.reserve(elements); in ExtractPrefix() 32 tuple_elements.push_back( in ExtractPrefix() 38 HloInstruction::CreateTuple(tuple_elements)); in ExtractPrefix() 48 std::vector<HloInstruction*> tuple_elements; in AppendSuffix() local 49 tuple_elements.reserve(input_shape.tuple_shapes_size()); in AppendSuffix() 51 tuple_elements.push_back( in AppendSuffix() 55 tuple_elements.insert(tuple_elements.end(), trailing_values.begin(), in AppendSuffix() 58 HloInstruction::CreateTuple(tuple_elements)); in AppendSuffix()
|
H A D | hlo_sharding_test.cc | 476 ASSERT_EQ(sharding_new_metadata.tuple_elements().size(), 3); in TEST_F() 478 ASSERT_EQ(sharding_new_metadata.tuple_elements()[0].metadata().size(), 1); in TEST_F() 480 sharding_new_metadata.tuple_elements()[0].metadata().front(), in TEST_F() 483 ASSERT_EQ(sharding_new_metadata.tuple_elements()[1].metadata().size(), 2); in TEST_F() 486 sharding_new_metadata.tuple_elements()[1].metadata()[i], in TEST_F() 490 ASSERT_EQ(sharding_new_metadata.tuple_elements()[2].metadata().size(), 1); in TEST_F() 492 sharding_new_metadata.tuple_elements()[2].metadata().front(), in TEST_F() 530 ASSERT_EQ(sharding_new_metadata.tuple_elements().size(), 3); in TEST_F() 532 for (const auto& sub_sharding : sharding_new_metadata.tuple_elements()) { in TEST_F() 566 EXPECT_EQ(sharding_no_metadata.tuple_elements().size(), 3); in TEST_F() [all …]
|
H A D | stable_sort_expander.cc | 113 std::vector<HloInstruction*> tuple_elements; in ExpandInstruction() local 114 tuple_elements.reserve(sort->operand_count()); in ExpandInstruction() 116 tuple_elements.push_back( in ExpandInstruction() 120 expanded_sort = tuple_elements[0]; in ExpandInstruction() 121 if (tuple_elements.size() > 1) { in ExpandInstruction() 123 HloInstruction::CreateTuple(tuple_elements)); in ExpandInstruction()
|
H A D | sharding_propagation_test.cc | 825 {copy_param0->sharding(), reduce->sharding().tuple_elements()[0], in TEST_P() 826 reduce->sharding().tuple_elements()[1]}) { in TEST_P() 880 EXPECT_THAT(gte1->sharding().tuple_elements()[0], in TEST_P() 882 EXPECT_THAT(gte1->sharding().tuple_elements()[1], in TEST_P() 887 {gte->sharding(), gte1->sharding().tuple_elements()[0], in TEST_P() 888 gte1->sharding().tuple_elements()[1], gte2->sharding()}) { in TEST_P() 936 EXPECT_THAT(tuple->sharding().tuple_elements()[0], in TEST_P() 938 EXPECT_THAT(tuple->sharding().tuple_elements()[1], ShardingMetadata({})); in TEST_P() 939 EXPECT_THAT(tuple1->sharding().tuple_elements()[0], ShardingMetadata({})); in TEST_P() 940 EXPECT_THAT(tuple1->sharding().tuple_elements()[1], in TEST_P() [all …]
|
H A D | sharding_propagation.cc | 61 return absl::c_any_of(sharding.tuple_elements(), IsSpatiallyPartitioned); in IsSpatiallyPartitioned() 788 for (int i = 0, e = a.tuple_elements().size(); i < e; ++i) { in SameShardingMetadata() 789 if (!same_metadata(a.tuple_elements()[i].metadata(), in SameShardingMetadata() 790 b.tuple_elements()[i].metadata())) { in SameShardingMetadata() 1297 sharding->tuple_elements(), in NormalizeDomain() 1487 new_sharding.tuple_elements()[sharding_index] = user.sharding(); in GetShardingFromUser() 1490 if (user.sharding().tuple_elements().empty()) { in GetShardingFromUser() 1499 user.sharding().tuple_elements()[0])); in GetShardingFromUser() 1500 for (int64_t i = 0; i < user.sharding().tuple_elements().size(); ++i) { in GetShardingFromUser() 1501 new_sharding.tuple_elements()[sharding_index + i] = in GetShardingFromUser() [all …]
|
H A D | hlo_element_type_converter.cc | 91 std::vector<HloInstruction*> tuple_elements; in ConvertTupleElements() local 102 tuple_elements.push_back(element); in ConvertTupleElements() 105 HloInstruction::CreateTuple(tuple_elements)); in ConvertTupleElements()
|
H A D | hlo_sharding.h | 300 std::vector<HloSharding>& tuple_elements() { return tuple_elements_; } in tuple_elements() function 301 const std::vector<HloSharding>& tuple_elements() const { in tuple_elements() function
|
H A D | hlo_sharding.cc | 376 for (auto& tuple_element_sharding : tuple_elements()) { in UsedDevices() 382 element_count = tuple_elements().size(); in UsedDevices() 852 for (HloSharding& sub_sharding : sharding.tuple_elements()) { in WithMetadata() 864 for (HloSharding& sub_sharding : sharding.tuple_elements()) { in WithoutMetadata()
|
H A D | hlo_sharding_util.cc | 49 const auto& lhs_shardings = lhs.tuple_elements(); in IsShardingMoreSpecific() 50 const auto& rhs_shardings = rhs.tuple_elements(); in IsShardingMoreSpecific() 83 for (int64_t i = 0; i < old.tuple_elements().size(); ++i) { in MergeSharding() 85 MergeSharding(old.tuple_elements()[i], &to_merge->tuple_elements()[i], in MergeSharding() 1216 for (const auto& subsharding : sharding.tuple_elements()) { in DevicesForShardingInternal()
|
H A D | hlo_instructions.cc | 1719 HloInstruction::InstructionVector tuple_elements; in CloneAndAppendInstructionIntoCalledComputation() local 1722 tuple_elements = root->operands(); in CloneAndAppendInstructionIntoCalledComputation() 1724 tuple_elements.push_back(root); in CloneAndAppendInstructionIntoCalledComputation() 1729 tuple_elements.push_back(inst); in CloneAndAppendInstructionIntoCalledComputation() 1732 tuple_elements.push_back(clone); in CloneAndAppendInstructionIntoCalledComputation() 1735 HloInstruction::CreateTuple(tuple_elements)); in CloneAndAppendInstructionIntoCalledComputation() 1750 int64_t index = tuple_elements.size(); in CloneAndAppendInstructionIntoCalledComputation()
|
H A D | hlo_sharding_metadata.cc | 243 domain_sharding.tuple_elements().size()); in ApplyShardingFromUsers()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/ |
H A D | shape_util.cc | 31 int64_t tuple_elements = xla::ShapeUtil::TupleElementCount(shape); in PopulateInfeedLayoutVector() local 32 for (int64_t i = 0; i < tuple_elements; ++i) { in PopulateInfeedLayoutVector() 177 int64_t tuple_elements = xla::ShapeUtil::TupleElementCount(input_shape); in GetShapeWithLayout() local 179 shapes.reserve(tuple_elements); in GetShapeWithLayout() 181 for (int64_t i = 0; i < tuple_elements; ++i) { in GetShapeWithLayout()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/tests/ |
H A D | transfer_manager_test.cc | 402 std::vector<Literal> tuple_elements; in Run() local 404 tuple_elements.push_back( in Run() 407 Literal literal = LiteralUtil::MakeTupleOwned(std::move(tuple_elements)); in Run() 431 std::vector<Literal> tuple_elements; in Run() local 433 tuple_elements.push_back( in Run() 436 Literal literal = LiteralUtil::MakeTupleOwned(std::move(tuple_elements)); in Run()
|
/aosp_15_r20/external/pytorch/c10/util/ |
H A D | Metaprogramming.h | 115 constexpr auto tuple_elements(Tuple t, std::index_sequence<Is...>) { 136 return tuple_elements(t, std::make_index_sequence<N>{}); 146 return tuple_elements(t, make_offset_index_sequence<size + N, -N>{}); 167 return tuple_elements(t, make_offset_index_sequence<Start, N>{});
|
/aosp_15_r20/external/pytorch/torch/csrc/jit/backends/ |
H A D | backend_detail.cpp | 341 auto tuple_elements = out_tuple_ty->elements(); in codegen_backend_module() local 342 type_check_ss << tuple_elements[0]->annotation_str() << ")"; in codegen_backend_module() 344 for (unsigned i = 1, e = tuple_elements.size(); i < e; ++i) { in codegen_backend_module() 349 << tuple_elements[i]->annotation_str() << ")"; in codegen_backend_module()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/gpu/ |
H A D | custom_call_test.cc | 232 std::vector<XlaOp> tuple_elements = BuildInputs(b, str); in BuildInputs() local 233 values.push_back(Tuple(&b, tuple_elements)); in BuildInputs() 251 std::vector<Shape> tuple_elements = BuildOutputType(str); in BuildOutputType() local 252 shapes.push_back(ShapeUtil::MakeTupleShape(tuple_elements)); in BuildOutputType()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/g3doc/ |
H A D | hlo_parser.md | 84 | '(' tuple_elements ')' 86 tuple_elements
|
/aosp_15_r20/external/pytorch/torch/testing/_internal/ |
H A D | common_jit.py | 319 tuple_elements = output.elements() 320 for i in range(len(tuple_elements)): 321 test_type(tuple_elements[i], out_sizes[i])
|
/aosp_15_r20/external/pytorch/c10/test/util/ |
H A D | Metaprogramming_test.cpp | 122 auto y = tuple_elements(x, std::index_sequence<0, 2>()); in TEST() 129 auto y = tuple_elements(x, std::index_sequence<0, 2, 1>()); in TEST()
|
/aosp_15_r20/external/pytorch/torch/csrc/jit/passes/ |
H A D | lower_tuples.cpp | 165 const auto& tuple_elements = tuple->elements(); in RemoveTupleConstants() local 168 for (const auto& elem : tuple_elements) { in RemoveTupleConstants()
|
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/ |
H A D | infeed_ops.cc | 134 int64_t tuple_elements = xla::ShapeUtil::TupleElementCount(*output_shape); in GetInfeedShapeWithLayout() local 135 for (int64_t i = 0; i < tuple_elements; ++i) { in GetInfeedShapeWithLayout()
|
/aosp_15_r20/external/pytorch/torch/_inductor/codegen/rocm/ |
H A D | ck_universal_gemm_template.py | 230 tuple_elements = ", ".join(map(str, iter(field_value)))
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/ |
H A D | literal_comparison.cc | 208 const int64_t tuple_elements = ShapeUtil::TupleElementCount(shape); in RecursiveElementCount() local 210 for (int64_t i = 0; i < tuple_elements; ++i) { in RecursiveElementCount()
|
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/spmd/ |
H A D | spmd_partitioner.cc | 2095 std::vector<HloSharding> subshardings = sharding.tuple_elements(); in Preprocess() 2121 hlo->sharding().tuple_elements(), in Preprocess() 2139 absl::c_any_of(hlo->sharding().tuple_elements(), in Preprocess() 2161 elements.reserve(sharding.tuple_elements().size()); in Preprocess() 2162 CHECK(!sharding.tuple_elements().empty()); in Preprocess() 2165 sharding.tuple_elements()[0]); in Preprocess() 2174 for (int64_t i = 1; i < sharding.tuple_elements().size(); ++i) { in Preprocess() 2177 sharding.tuple_elements()[i]), in Preprocess() 4633 TF_RET_CHECK(absl::c_all_of(root_sharding.tuple_elements(), in PreprocessSharding()
|
H A D | custom_call_handler.cc | 430 hlo->sharding().tuple_elements().size(), in HandleCustomCall()
|