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

Merge pull request #25237 from jakevdp:faster-isscalar

Browse Source 
PiperOrigin-RevId: 702517550
This commit is contained in:
jax authors 2024-12-03 17:03:10 -08:00
commit 40122f7c03
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
jax/_src/numpy/lax_numpy.py
View File

@ -624,9 +624,11 @@ def isscalar(element: Any) -> bool:
>>> jnp.isscalar(slice(10))
False
"""
if (isinstance(element, (np.ndarray, jax.Array))
or hasattr(element, '__jax_array__')
or np.isscalar(element)):
if np.isscalar(element):
return True
elif isinstance(element, (np.ndarray, jax.Array)):
return element.ndim == 0
elif hasattr(element, '__jax_array__'):
return asarray(element).ndim == 0
return False