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

Fix cuda array interface with old jaxlib.

Browse Source 
arg name was added in xla_extension_version 261.

PiperOrigin-RevId: 629745255
This commit is contained in:
Jieying Luo 2024-05-01 09:25:46 -07:00 committed by jax authors
parent e75e4a5991
commit 26049b1059
1 changed files with 1 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -2538,9 +2538,7 @@ def array(object: Any, dtype: DTypeLike | None = None, copy: bool = True,
cai=cai, gpu_backend=backend, device_id=device_id
)
else:
object = xc._xla.cuda_array_interface_to_buffer(
cai=cai, gpu_backend=backend
)
object = xc._xla.cuda_array_interface_to_buffer(cai, backend)
object = tree_map(lambda leaf: leaf.__jax_array__()
if hasattr(leaf, "__jax_array__") else leaf, object)