xax 0.2.15__tar.gz → 0.2.16__tar.gz

{xax-0.2.15/xax.egg-info → xax-0.2.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.15
+Version: 0.2.16
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.2.15 → xax-0.2.16}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.15"
+__version__ = "0.2.16"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

{xax-0.2.15 → xax-0.2.16}/xax/nn/geom.py

@@ -102,12 +102,13 @@ def get_projected_gravity_vector_from_quat(quat: Array, eps: float = 1e-6) -> Ar
     return jnp.concatenate([gx, gy, -gz], axis=-1)
 
 
-def rotate_vector_by_quat(vector: Array, quat: Array, eps: float = 1e-6) -> Array:
+def rotate_vector_by_quat(vector: Array, quat: Array, inverse: bool = False, eps: float = 1e-6) -> Array:
     """Rotates a vector by a quaternion.
 
     Args:
         vector: The vector to rotate, shape (*, 3).
         quat: The quaternion to rotate by, shape (*, 4).
+        inverse: If True, rotate the vector by the conjugate of the quaternion.
         eps: A small epsilon value to avoid division by zero.
 
     Returns:
@@ -117,6 +118,9 @@ def rotate_vector_by_quat(vector: Array, quat: Array, eps: float = 1e-6) -> Arra
     quat = quat / (jnp.linalg.norm(quat, axis=-1, keepdims=True) + eps)
     w, x, y, z = jnp.split(quat, 4, axis=-1)
 
+    if inverse:
+        x, y, z = -x, -y, -z
+
     # Extract vector components
     vx, vy, vz = jnp.split(vector, 3, axis=-1)
 

{xax-0.2.15 → xax-0.2.16/xax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.15
+Version: 0.2.16
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte