Searched refs:tuple_elements (Results 1 – 12 of 12) sorted by relevance
29 std::vector<HloInstruction*> tuple_elements; in ExtractPrefix() local30 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() local49 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()
113 std::vector<HloInstruction*> tuple_elements; in ExpandInstruction() local114 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()
89 std::vector<HloInstruction*> tuple_elements; in ConvertTupleElements() local100 tuple_elements.push_back(element); in ConvertTupleElements()103 HloInstruction::CreateTuple(tuple_elements)); in ConvertTupleElements()
210 const std::vector<HloSharding>& tuple_elements() const { in tuple_elements() function
122 for (auto& tuple_element_sharding : tuple_elements()) { in UsedDevices()128 element_count = tuple_elements().size(); in UsedDevices()
1476 HloInstruction::InstructionVector tuple_elements; in CloneAndFuseInternal() local1479 tuple_elements = fused_root->operands(); in CloneAndFuseInternal()1481 tuple_elements.push_back(fused_root); in CloneAndFuseInternal()1486 tuple_elements.push_back(inst); in CloneAndFuseInternal()1489 tuple_elements.push_back(clone); in CloneAndFuseInternal()1492 HloInstruction::CreateTuple(tuple_elements)); in CloneAndFuseInternal()1507 int64 index = tuple_elements.size(); in CloneAndFuseInternal()
239 domain_sharding.tuple_elements().size()); in ApplyShardingFromUsers()
31 int64 tuple_elements = xla::ShapeUtil::TupleElementCount(shape); in PopulateInfeedLayoutVector() local32 for (int64 i = 0; i < tuple_elements; ++i) { in PopulateInfeedLayoutVector()135 int64 tuple_elements = xla::ShapeUtil::TupleElementCount(input_shape); in GetShapeWithLayout() local137 shapes.reserve(tuple_elements); in GetShapeWithLayout()139 for (int64 i = 0; i < tuple_elements; ++i) { in GetShapeWithLayout()
322 std::vector<Literal> tuple_elements; in Run() local324 tuple_elements.push_back( in Run()327 Literal literal = LiteralUtil::MakeTupleOwned(std::move(tuple_elements)); in Run()353 std::vector<Literal> tuple_elements; in Run() local355 tuple_elements.push_back( in Run()358 Literal literal = LiteralUtil::MakeTupleOwned(std::move(tuple_elements)); in Run()
84 | '(' tuple_elements ')'86 tuple_elements
190 const int64 tuple_elements = ShapeUtil::TupleElementCount(shape); in RecursiveElementCount() local192 for (int64 i = 0; i < tuple_elements; ++i) { in RecursiveElementCount()
1662 Literal tuple_elements[] = { in TEST_F() local1668 {&tuple_elements[0], &tuple_elements[1], &nil_literal}); in TEST_F()