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

Explicitly check for DShapedArray in _to_logical_sharding instead of returning a sharding by default.

Browse Source
This commit is contained in:
Alexey Radul 2023-07-12 13:08:57 -04:00
parent ef9f1cbec3
commit 4794a827ca
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -49,6 +49,7 @@ from jax._src import util
from jax._src import xla_bridge as xb
from jax._src.abstract_arrays import array_types
from jax._src.config import config
from jax._src.core import DShapedArray
from jax._src.core import ShapedArray
from jax._src.interpreters import ad
from jax._src.interpreters import batching
@ -2125,14 +2126,13 @@ def _to_logical_sharding(
) -> Optional[sharding_impls.XLACompatibleSharding]:
if is_unspecified(sharding) or is_auto(sharding):
return None
elif isinstance(aval, ShapedArray):
elif isinstance(aval, (ShapedArray, DShapedArray)):
assert isinstance(sharding, sharding_impls.XLACompatibleSharding)
return sharding
elif isinstance(aval, core.AbstractToken):
return None
else:
assert isinstance(sharding, sharding_impls.XLACompatibleSharding)
return sharding
raise TypeError(aval)
@profiler.annotate_function