Searched refs:lhs_sharding (Results 1 – 2 of 2) sorted by relevance
393 const HloSharding& lhs_sharding, const HloSharding& rhs_sharding) { in GetWindowedEinsumConfiguration() argument395 output_sharding_transposed_to_match_lhs == lhs_sharding && in GetWindowedEinsumConfiguration()484 const HloSharding& lhs_sharding = lhs.sharding(); in PartitionBaseCase() local486 if (lhs_sharding.ReplicateOnLastTileDim() || in PartitionBaseCase()496 lhs_sharding, indices_map.lhs_to_rhs_indices, in PartitionBaseCase()504 lhs_sharding, indices_map.lhs_to_output_indices, in PartitionBaseCase()1479 output_sharding_transposed_to_match_rhs, lhs_sharding, rhs_sharding); in PartitionBaseCase()2098 auto lhs_sharding = lhs.sharding(); in PartitionDotGroupOnContracting() local2100 auto lhs_tile_shape = lhs_sharding.tile_assignment().dimensions(); in PartitionDotGroupOnContracting()2108 group_count *= lhs_sharding.tile_assignment().dim(dim.lhs); in PartitionDotGroupOnContracting()[all …]
258 auto unsupported_sharding = [&](const HloSharding& lhs_sharding, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()262 return lhs_sharding.tile_assignment().dim(dnums.input_batch_dimension()) != in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()569 auto unsupported_sharding = [&](const HloSharding& lhs_sharding, in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()571 return lhs_sharding.tile_assignment().dim(dnums.input_batch_dimension()) != in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()