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

Merge pull request #7943 from rsepassi:main

Browse Source 
PiperOrigin-RevId: 397871285
This commit is contained in:
jax authors 2021-09-20 16:42:00 -07:00
commit 2d03bde7e5
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -301,7 +301,7 @@ def _result_dtype(op, *args):
def _arraylike(x):
return isinstance(x, ndarray) or isscalar(x) or hasattr(x, '__jax_array__')
return isinstance(x, ndarray) or hasattr(x, '__jax_array__') or isscalar(x)
def _check_arraylike(fun_name, *args):
"""Check if all args fit JAX's definition of arraylike."""