Searched refs:sub_shardings (Results 1 – 3 of 3) sorted by relevance
108 HloSharding HloSharding::Tuple(const ShapeTree<HloSharding>& sub_shardings) { in Tuple() argument110 flattened_list.reserve(sub_shardings.leaf_count()); in Tuple()111 for (const auto& index_to_sharding : sub_shardings.leaves()) { in Tuple()120 flattened_list.push_back(sub_shardings.element(ShapeIndex({}))); in Tuple()639 std::vector<HloSharding> sub_shardings( in GetSubSharding() local641 return HloSharding::Tuple(*sub_shape, sub_shardings); in GetSubSharding()
650 std::vector<HloSharding> sub_shardings = in InferShardingFromOperands() local661 sub_shardings[sub_sharding_index + i])) { in InferShardingFromOperands()662 sub_shardings[sub_sharding_index + i] = in InferShardingFromOperands()668 operand->sharding(), sub_shardings[sub_sharding_index])) { in InferShardingFromOperands()669 sub_shardings[sub_sharding_index] = operand->sharding(); in InferShardingFromOperands()676 HloSharding new_sharding = HloSharding::Tuple(shape, sub_shardings); in InferShardingFromOperands()
89 static HloSharding Tuple(const ShapeTree<HloSharding>& sub_shardings);