Searched refs:other_sharding (Results 1 – 1 of 1) sorted by relevance
1916 const HloSharding& other_sharding, const HloSharding& output_sharding, in GetNonContractingPartitionGroupedShardingForOtherOperand() argument1930 if (other_sharding.ReplicateOnLastTileDim() && in GetNonContractingPartitionGroupedShardingForOtherOperand()1931 other_sharding.tile_assignment().dimensions().back() % group_count == 0) { in GetNonContractingPartitionGroupedShardingForOtherOperand()1933 other_sharding.tile_assignment().num_dimensions() - 1); in GetNonContractingPartitionGroupedShardingForOtherOperand()1943 other_sharding, output_grouped.device_groups)) { in GetNonContractingPartitionGroupedShardingForOtherOperand()1963 other_sharding.tile_assignment().num_dimensions() - 1) { in GetNonContractingPartitionGroupedShardingForOtherOperand()1966 other_sharding, {other_group_dims[0]}, in GetNonContractingPartitionGroupedShardingForOtherOperand()1967 {other_sharding.tile_assignment().dimensions().back() / in GetNonContractingPartitionGroupedShardingForOtherOperand()1971 } else if (!other_sharding.IsReplicated()) { in GetNonContractingPartitionGroupedShardingForOtherOperand()1973 GroupShardingOnDims(other_sharding, other_group_dims), output_grouped, in GetNonContractingPartitionGroupedShardingForOtherOperand()[all …]