Home
last modified time | relevance | path

Searched refs:tuple_elements (Results 1 – 25 of 28) sorted by relevance

12

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/
H A Dtuple_util.cc29 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 Dhlo_sharding_test.cc476 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 Dstable_sort_expander.cc113 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 Dsharding_propagation_test.cc825 {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 Dsharding_propagation.cc61 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 Dhlo_element_type_converter.cc91 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 Dhlo_sharding.h300 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 Dhlo_sharding.cc376 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 Dhlo_sharding_util.cc49 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 Dhlo_instructions.cc1719 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 Dhlo_sharding_metadata.cc243 domain_sharding.tuple_elements().size()); in ApplyShardingFromUsers()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/tf2xla/
H A Dshape_util.cc31 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 Dtransfer_manager_test.cc402 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 DMetaprogramming.h115 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 Dbackend_detail.cpp341 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 Dcustom_call_test.cc232 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 Dhlo_parser.md84 | '(' tuple_elements ')'
86 tuple_elements
/aosp_15_r20/external/pytorch/torch/testing/_internal/
H A Dcommon_jit.py319 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 DMetaprogramming_test.cpp122 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 Dlower_tuples.cpp165 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 Dinfeed_ops.cc134 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 Dck_universal_gemm_template.py230 tuple_elements = ", ".join(map(str, iter(field_value)))
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/
H A Dliteral_comparison.cc208 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 Dspmd_partitioner.cc2095 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 Dcustom_call_handler.cc430 hlo->sharding().tuple_elements().size(), in HandleCustomCall()

12