Searched refs:rhs_sharding (Results 1 – 2 of 2) sorted by relevance
393 const HloSharding& lhs_sharding, const HloSharding& rhs_sharding) { in GetWindowedEinsumConfiguration() argument420 output_sharding_transposed_to_match_rhs == rhs_sharding && in GetWindowedEinsumConfiguration()485 const HloSharding& rhs_sharding = rhs.sharding(); in PartitionBaseCase() local487 rhs_sharding.ReplicateOnLastTileDim() || in PartitionBaseCase()500 rhs_sharding, indices_map.rhs_to_lhs_indices, in PartitionBaseCase()508 rhs_sharding, indices_map.rhs_to_output_indices, in PartitionBaseCase()523 lhs_sharding_transposed_to_match_rhs == rhs_sharding) { in PartitionBaseCase()1479 output_sharding_transposed_to_match_rhs, lhs_sharding, rhs_sharding); in PartitionBaseCase()2099 auto rhs_sharding = rhs.sharding(); in PartitionDotGroupOnContracting() local2101 auto rhs_tile_shape = rhs_sharding.tile_assignment().dimensions(); in PartitionDotGroupOnContracting()[all …]
259 const HloSharding& rhs_sharding) { in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS() argument264 rhs_sharding.tile_assignment().dim( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnRHS()570 const HloSharding& rhs_sharding) { in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS() argument573 rhs_sharding.tile_assignment().dim( in PartitionConvolutionWithSpatialDimensionHaloExchangeOnLHS()