Searched refs:UngroupSharding (Results 1 – 3 of 3) sorted by relevance
1700 rhs.sharding() == UngroupSharding(AlignGroupsWith( in PartitionDotGroupOnBatch()1713 rhs = rhs.Reshard(UngroupSharding(rhs_grouped)); in PartitionDotGroupOnBatch()1716 lhs = lhs.Reshard(UngroupSharding(lhs_grouped)); in PartitionDotGroupOnBatch()1814 operand.sharding() != UngroupSharding(grouped)) { in PartitionDotGroupOnBatch()1817 auto resharded = operand.Reshard(UngroupSharding(grouped)); in PartitionDotGroupOnBatch()1874 dot->set_sharding(UngroupSharding(output_grouped)); in PartitionDotGroupOnBatch()2015 matching.sharding() != UngroupSharding(matching_grouped)) { in PartitionDotGroupOnNonContracting()2032 matching = matching.Reshard(UngroupSharding(matching_grouped)); in PartitionDotGroupOnNonContracting()2046 other = other.Reshard(UngroupSharding(*other_grouped)); in PartitionDotGroupOnNonContracting()2051 other = other.Reshard(UngroupSharding(*other_grouped)); in PartitionDotGroupOnNonContracting()[all …]
324 HloSharding UngroupSharding(const GroupedSharding& grouped_sharding);
1435 HloSharding UngroupSharding(const GroupedSharding& grouped_sharding) { in UngroupSharding() function1553 return UngroupSharding(AlignGroupsWith(sharding_grouped, reference_grouped)); in AlignShardingOnDims()