Searched refs:require_matching_devices_to_group (Results 1 – 1 of 1) sorted by relevance
2015 bool require_matching_devices_to_group, in PartitionDotGroupOnBatch() argument2056 if (require_matching_devices_to_group && lhs_rhs_dims_matching) { in PartitionDotGroupOnBatch()2176 if (require_matching_devices_to_group && in PartitionDotGroupOnBatch()2359 bool require_matching_devices_to_group, in PartitionDotGroupOnNonContracting() argument2383 if (require_matching_devices_to_group && in PartitionDotGroupOnNonContracting()2588 bool require_matching_devices_to_group, in PartitionDotGroupOnContracting() argument2621 if (require_matching_devices_to_group && rhs.sharding() != rhs_sharding) { in PartitionDotGroupOnContracting()2628 if (require_matching_devices_to_group && lhs.sharding() != lhs_sharding) { in PartitionDotGroupOnContracting()2996 bool require_matching_devices_to_group, SpmdBuilder* b, in PrioritizeContractingDimensionsPartitioning() argument3062 if (require_matching_devices_to_group && rhs.sharding() != rhs_sharding) { in PrioritizeContractingDimensionsPartitioning()[all …]