Fix deprecations according to JAX 0.4.17

flferretti · Oct 4, 2023 · 1f84317 · 1f84317
1 parent 48c54f0
commit 1f84317
Show file tree

Hide file tree

Showing 2 changed files with 0 additions and 4 deletions.
diff --git a/src/jaxsim/utils/jaxsim_dataclass.py b/src/jaxsim/utils/jaxsim_dataclass.py
@@ -4,9 +4,7 @@
 import dataclasses
 from typing import Generator
 
-import jax.abstract_arrays
 import jax.flatten_util
-import jax.interpreters.partial_eval
 import jax_dataclasses
 
 import jaxsim.typing as jtp

diff --git a/src/jaxsim/utils/tracing.py b/src/jaxsim/utils/tracing.py
@@ -1,7 +1,6 @@
 from typing import Any
 
 import jax._src.core
-import jax.abstract_arrays
 import jax.flatten_util
 import jax.interpreters.partial_eval
 
@@ -14,7 +13,6 @@ def tracing(var: Any) -> bool | jax.Array:
             isinstance(var, t)
             for t in (
                 jax._src.core.Tracer,
-                jax.abstract_arrays.ShapedArray,
                 jax.interpreters.partial_eval.DynamicJaxprTracer,
             )
         ]