mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Don't apply shardy config during def_partition.

Browse Source 
PR #25834 intended to dynamically choose the the partitioner API, but
it still applies the configuration value too early (it should only be
applied in __call__, not in def_partition and __call__).
This commit is contained in:
Johannes Reifferscheid 2025-01-28 15:37:47 +01:00
parent e332b94f19
commit 55d891c5bf
1 changed files with 0 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jax/_src/custom_partitioning.py
View File

@ -454,11 +454,6 @@ class custom_partitioning:
def def_partition(self, partition, infer_sharding_from_operands,
propagate_user_sharding=None, decode_shardings=True,
sharding_rule=None):
if config.use_shardy_partitioner.value:
infer_sharding_from_operands = None
propagate_user_sharding = None
else:
sharding_rule = None
self.partition = partition
self.propagate_user_sharding = propagate_user_sharding
self.infer_sharding_from_operands = infer_sharding_from_operands