xax 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

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/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.dist-info → xax-0.2.17.dist-info}/METADATA

@@ -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.dist-info → xax-0.2.17.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=JVxuGfbwBPHXiF4kSG0Pb73mzu3EIaRipjvt0Y-Z9W4,15733
+xax/__init__.py,sha256=dBxZ_r1ck3C9ZH9VRM38-ApVkyUe7CnI_0SF9k07KcI,15733
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
@@ -10,7 +10,7 @@ xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/equinox.py,sha256=JZuSApD4bL0UK5W1nrQtucWYvNWUha07J6LTLk_RX-Y,4910
 xax/nn/export.py,sha256=pRfM2B4hB2EvljysC6AjtgB_7Cn7JtaP3dhYU2stZtY,5545
 xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
-xax/nn/geom.py,sha256=B8QE-L-xJWhf9KygTByPUAWe7Clpek4GlTABpsJFMBs,7702
+xax/nn/geom.py,sha256=A7WPefMvgwUNReZC7_HX1GmvHPASyghbaXaKsuhwDrE,7382
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
 xax/nn/metrics.py,sha256=OAkeScwhi-wTBIJ59KHUhYbZTq4V4V-LG-mKlxMJ7bY,3238
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
@@ -58,8 +58,8 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.15.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.15.dist-info/METADATA,sha256=6LJoiKOyNmF1MJSwVSdbEJATzSv1P77Amn4ZJCbWaP0,1880
-xax-0.2.15.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-xax-0.2.15.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.15.dist-info/RECORD,,
+xax-0.2.17.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.17.dist-info/METADATA,sha256=GQhyzReeHSrZkYrpxeSXt19z2271zD49-S6fwN6cagU,1880
+xax-0.2.17.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+xax-0.2.17.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.17.dist-info/RECORD,,

{xax-0.2.15.dist-info → xax-0.2.17.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (79.0.0)
+Generator: setuptools (80.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any