Fixes for jax_dataclasses API change

brentyi · Jun 2, 2021 · 7cd3e37 · 7cd3e37
1 parent 7f45ef3
commit 7cd3e37
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 5 deletions.
diff --git a/jaxlie/_se2.py b/jaxlie/_se2.py
@@ -15,7 +15,7 @@
     tangent_dim=3,
     space_dim=2,
 )
-@jax_dataclasses.dataclass
+@jax_dataclasses.pytree_dataclass
 class SE2(_base.SEBase[SO2]):
     """Special Euclidean group for proper rigid transforms in 2D."""
 

diff --git a/jaxlie/_se3.py b/jaxlie/_se3.py
@@ -28,7 +28,7 @@ def _skew(omega: hints.Vector) -> hints.MatrixJax:
     tangent_dim=6,
     space_dim=3,
 )
-@jax_dataclasses.dataclass
+@jax_dataclasses.pytree_dataclass
 class SE3(_base.SEBase[SO3]):
     """Special Euclidean group for proper rigid transforms in 3D."""
 

diff --git a/jaxlie/_so2.py b/jaxlie/_so2.py
@@ -14,7 +14,7 @@
     tangent_dim=1,
     space_dim=2,
 )
-@jax_dataclasses.dataclass
+@jax_dataclasses.pytree_dataclass
 class SO2(_base.SOBase):
     """Special orthogonal group for 2D rotations."""
 

diff --git a/jaxlie/_so3.py b/jaxlie/_so3.py
@@ -14,7 +14,7 @@
     tangent_dim=3,
     space_dim=3,
 )
-@jax_dataclasses.dataclass
+@jax_dataclasses.pytree_dataclass
 class SO3(_base.SOBase):
     """Special orthogonal group for 3D rotations."""
 

diff --git a/setup.py b/setup.py
@@ -21,7 +21,7 @@
         "jax",
         "jaxlib",
         "numpy",
-        "jax_dataclasses",
+        "jax_dataclasses>=1.0.0",
         # `overrides` should not be updated until the following issues are resolved:
         # > https://github.com/mkorpela/overrides/issues/65
         # > https://github.com/mkorpela/overrides/issues/63