Searched refs:split_dims_added (Results 1 – 1 of 1) sorted by relevance
3572 int64 split_dims_added = 0; in AllGatherShardsInternal() local3574 if (sharding.tile_assignment().dim(i - split_dims_added) == 1 || in AllGatherShardsInternal()3575 !absl::c_linear_search(selected_dims, i - split_dims_added)) { in AllGatherShardsInternal()3576 xpose_permutation[i] = i + tiled_dims.size() - split_dims_added; in AllGatherShardsInternal()3578 xpose_permutation[i] = split_dims_added; in AllGatherShardsInternal()3579 xpose_permutation[i + 1] = i + tiled_dims.size() - split_dims_added; in AllGatherShardsInternal()3580 split_dims_added++; in AllGatherShardsInternal()