fix(integrations): fix type check for jax.Array introduced in jax==0.…

…4.1 (#4718) * fix(integrations): fix type check for jax.Array introduced in jax==0.4.1
wandb · Jan 4, 2023 · 3c04cf5 · 3c04cf5
1 parent 4f272d3
commit 3c04cf5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/requirements_dev.txt b/requirements_dev.txt
@@ -28,7 +28,7 @@ docker
 stable_baselines3
 tensorboard
 gym
-jax[cpu]<0.4; sys_platform == 'darwin' or sys_platform == 'linux'
+jax[cpu]; sys_platform == 'darwin' or sys_platform == 'linux'
 fastcore; python_version > '3.6'
 fastcore==1.3.29; python_version == '3.6'
 pyarrow

diff --git a/wandb/util.py b/wandb/util.py
@@ -542,7 +542,7 @@ def is_pytorch_tensor_typename(typename: str) -> bool:
 
 
 def is_jax_tensor_typename(typename: str) -> bool:
-    return typename.startswith("jaxlib.") and "DeviceArray" in typename
+    return typename.startswith("jaxlib.") and "Array" in typename
 
 
 def get_jax_tensor(obj: Any) -> Optional[Any]: