Searched refs:grouped_sharding (Results 1 – 2 of 2) sorted by relevance
1435 HloSharding UngroupSharding(const GroupedSharding& grouped_sharding) { in UngroupSharding() argument1438 auto grouped_tiling = grouped_sharding.sharding.tile_assignment(); in UngroupSharding()1439 if (grouped_sharding.sharding.IsTileMaximal()) { in UngroupSharding()1440 tiling_dims = std::vector<int64>(grouped_sharding.data_rank, 1); in UngroupSharding()1441 if (grouped_sharding.device_groups[0].size() != 1) { in UngroupSharding()1443 tiling_dims.push_back(grouped_sharding.device_groups[0].size()); in UngroupSharding()1449 partial_sharding = grouped_sharding.sharding.ReplicateOnLastTileDim(); in UngroupSharding()1450 tiling_dims = grouped_sharding.sharding.tile_assignment().dimensions(); in UngroupSharding()1451 if (absl::c_linear_search(grouped_sharding.group_dims, in UngroupSharding()1458 for (int64 i = 0; i < grouped_sharding.group_dims.size(); ++i) { in UngroupSharding()[all …]
301 HloSharding grouped_sharding) in GroupedSharding()306 sharding(std::move(grouped_sharding)) {} in GroupedSharding()324 HloSharding UngroupSharding(const GroupedSharding& grouped_sharding);328 GroupedSharding AlignGroupsWith(GroupedSharding grouped_sharding,342 Shape GetPerGroupBaseShape(const GroupedSharding& grouped_sharding,