xax 0.2.15__tar.gz → 0.2.17__tar.gz

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

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

{xax-0.2.15 → xax-0.2.17}/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.17"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

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

@@ -86,28 +86,16 @@ def get_projected_gravity_vector_from_quat(quat: Array, eps: float = 1e-6) -> Ar
     Returns:
         A 3D vector representing the gravity in the local frame, shape (*, 3).
     """
-    # Normalize quaternion
-    quat = quat / (jnp.linalg.norm(quat, axis=-1, keepdims=True) + eps)
-    w, x, y, z = jnp.split(quat, 4, axis=-1)
-
-    # Gravity vector in world frame is [0, 0, -1] (pointing down)
-    # Rotate gravity vector using quaternion rotation
+    return rotate_vector_by_quat(jnp.array([0, 0, -9.81]), quat, inverse=True, eps=eps)
 
-    # Calculate quaternion rotation: q * [0,0,-1] * q^-1
-    gx = 2 * (x * z - w * y)
-    gy = 2 * (y * z + w * x)
-    gz = w * w - x * x - y * y + z * z
 
-    # Note: We're rotating [0,0,-1], so we negate gz to match the expected direction
-    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 +105,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.17/xax.egg-info}/PKG-INFO

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