Fix get_batch_axes()

brentyi · Feb 1, 2023 · ad93513 · ad93513
1 parent 76c3042
commit ad93513
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/jaxlie/utils/_utils.py b/jaxlie/utils/_utils.py
@@ -46,7 +46,9 @@ def _wrap(cls: Type[T]) -> Type[T]:
 
         # JIT all methods.
         for f in filter(
-            lambda f: not f.startswith("_") and callable(getattr(cls, f)),
+            lambda f: not f.startswith("_")
+            and callable(getattr(cls, f))
+            and f != "get_batch_axes",  # Avoid returning tracers.
             dir(cls),
         ):
             setattr(cls, f, jax.jit(getattr(cls, f)))

diff --git a/setup.py b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="jaxlie",
-    version="1.3.0",
+    version="1.3.1",
     description="Matrix Lie groups in Jax",
     long_description=long_description,
     long_description_content_type="text/markdown",