Searched refs:xpose_permutation (Results 1 – 1 of 1) sorted by relevance
3571 std::vector<int64> xpose_permutation(result->shape().rank()); in AllGatherShardsInternal() local3573 for (int64 i = 0; i < xpose_permutation.size(); ++i) { 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()3585 ShapeInference::InferTransposeShape(result->shape(), xpose_permutation) in AllGatherShardsInternal()3587 result, xpose_permutation)); in AllGatherShardsInternal()