Searched refs:group_dim_shards (Results 1 – 2 of 2) sorted by relevance
1380 std::vector<int64> group_dim_shards(group_dims.size(), 1); in GroupShardingOnDims() local1381 return GroupShardingOnDims(sharding, group_dims, group_dim_shards); in GroupShardingOnDims()1386 absl::Span<const int64> group_dim_shards) { in GroupShardingOnDims() argument1392 CHECK_EQ(grouped_tiling_dims[group_dims[i]] % group_dim_shards[i], 0); in GroupShardingOnDims()1394 grouped_tiling_dims[group_dims[i]] / group_dim_shards[i]; in GroupShardingOnDims()1395 grouped_tiling_dims[group_dims[i]] = group_dim_shards[i]; in GroupShardingOnDims()1404 group_dim_shards[i]; in GroupShardingOnDims()1405 group_id += indices[group_dims[i]] / group_dim_shards[i]; in GroupShardingOnDims()
317 absl::Span<const int64> group_dim_shards);