Searched refs:aligned_rhs_sharding (Results 1 – 1 of 1) sorted by relevance
96 auto aligned_rhs_sharding = in PartitionConvolutionWithBatchGroupCount() local114 rhs = rhs.Reshard(aligned_rhs_sharding); in PartitionConvolutionWithBatchGroupCount()187 auto aligned_rhs_sharding = in PartitionConvolutionWithFeatureGroupCount() local206 rhs = rhs.Reshard(aligned_rhs_sharding); in PartitionConvolutionWithFeatureGroupCount()253 auto aligned_rhs_sharding = in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() local277 if (unsupported_sharding(lhs.sharding(), aligned_rhs_sharding)) { in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()281 rhs = rhs.Reshard(aligned_rhs_sharding).PadWithValue(zero); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()562 auto aligned_rhs_sharding = hlo_sharding_util::ReverseSharding( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() local586 if (unsupported_sharding(lhs.sharding(), aligned_rhs_sharding)) { in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()591 rhs.Reshard(aligned_rhs_sharding).PadWithValue(zero, reversed_rhs_dims); in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()