Searched refs:reversed_rhs_dims (Results 1 – 1 of 1) sorted by relevance
545 std::vector<int64> reversed_rhs_dims; in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() local548 reversed_rhs_dims.push_back(dnums.kernel_spatial_dimensions(i)); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()551 if (!reversed_rhs_dims.empty()) { in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()553 auto left_padded_rhs = HaloExchangeToPadOnLeft(rhs, reversed_rhs_dims); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()564 reversed_rhs_dims); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()566 hlo_sharding_util::ReverseSharding(rhs.sharding(), reversed_rhs_dims), in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()584 rhs = rhs.PadWithValue(zero, reversed_rhs_dims); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()591 rhs.Reshard(aligned_rhs_sharding).PadWithValue(zero, reversed_rhs_dims); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()