Searched refs:grouped_tiling (Results 1 – 1 of 1) sorted by relevance
1425 Array<int64> grouped_tiling(grouped_tiling_dims); in GroupShardingOnDims() local1426 grouped_tiling.FillIota(0); in GroupShardingOnDims()1430 ? HloSharding::PartialTile(grouped_tiling) in GroupShardingOnDims()1431 : HloSharding::Tile(grouped_tiling); in GroupShardingOnDims()1438 auto grouped_tiling = grouped_sharding.sharding.tile_assignment(); in UngroupSharding() local1446 grouped_tiling = Array<int64>(tiling_dims); in UngroupSharding()1447 grouped_tiling.FillIota(0); in UngroupSharding()1454 grouped_tiling.Reshape(tiling_dims); in UngroupSharding()1463 grouped_tiling.Each([&](absl::Span<const int64> indices, int64 device) { in UngroupSharding()1471 grouped_tiling.dim(dim) + in UngroupSharding()