Remove in-place operations, easier to port to numpy

brentyi · Apr 26, 2023 · 8d0c242 · 8d0c242
1 parent 654d631
commit 8d0c242
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/jaxlie/_so2.py b/jaxlie/_so2.py
@@ -69,7 +69,7 @@ def as_matrix(self) -> jax.Array:
         out = jnp.array(
             [
                 # [cos, -sin],
-                cos_sin.at[1].multiply(-1),
+                cos_sin * jnp.array([1, -1]),
                 # [sin, cos],
                 cos_sin[::-1],
             ]

diff --git a/jaxlie/_so3.py b/jaxlie/_so3.py
@@ -290,7 +290,7 @@ def apply(self, target: hints.Array) -> jax.Array:
         assert target.shape == (3,)
 
         # Compute using quaternion multiplys.
-        padded_target = jnp.zeros(4).at[1:].set(target)
+        padded_target = jnp.concatenate([jnp.zeros(1), target])
         return (self @ SO3(wxyz=padded_target) @ self.inverse()).wxyz[1:]
 
     @override