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

Resolve a TODO now that in_shardings are chosen by XLA for inputs that don't have sharding specified or are uncommitted

Browse Source 
PiperOrigin-RevId: 621991853
This commit is contained in:
Yash Katariya 2024-04-04 15:38:49 -07:00 committed by jax authors
parent 033992867f
commit b322d399e1
1 changed files with 10 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jax/_src/interpreters/pxla.py
View File

@ -2757,11 +2757,6 @@ def _maybe_get_and_check_in_shardings(
dtypes.issubdtype(aval.dtype, dtypes.extended)):
xla_s = aval.dtype._rules.logical_sharding(aval, xla_s)
new_in_shardings.append(xla_s)
else:
# TODO(yashkatariya): Remove the if branch for abstract_token once
# choosing input shardings by XLA is enabled again.
if aval is core.abstract_token:
new_in_shardings.append(orig)
else:
xla_hlo_s = xla_s._to_xla_hlo_sharding(aval.ndim) # type: ignore
orig_hlo_s = orig._to_xla_hlo_sharding(aval.ndim) # type: ignore