xax 0.2.23__py3-none-any.whl → 0.3.1__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.23"
+__version__ = "0.3.1"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -44,6 +44,7 @@ __all__ = [
     "normalize",
     "quat_to_euler",
     "quat_to_rotmat",
+    "quat_mul",
     "rotate_vector_by_quat",
     "rotation6d_to_rotation_matrix",
     "rotation_matrix_to_rotation6d",
@@ -225,6 +226,7 @@ NAME_MAP: dict[str, str] = {
     "normalize": "nn.geom",
     "quat_to_euler": "nn.geom",
     "quat_to_rotmat": "nn.geom",
+    "quat_mul": "nn.geom",
     "rotate_vector_by_quat": "nn.geom",
     "rotation6d_to_rotation_matrix": "nn.geom",
     "rotation_matrix_to_rotation6d": "nn.geom",
@@ -398,6 +400,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
         normalize,
+        quat_mul,
         quat_to_euler,
         quat_to_rotmat,
         rotate_vector_by_quat,

xax/nn/geom.py

@@ -251,3 +251,24 @@ def rotation_matrix_to_rotation6d(rotation_matrix: jnp.ndarray) -> jnp.ndarray:
     # Simply concatenate a1 and a2 from SO(3)
     r6d = jnp.concatenate([rotation_matrix[..., 0], rotation_matrix[..., 1]], axis=-1)
     return r6d.reshape(shape[:-2] + (6,))
+
+
+def quat_mul(q2: Array, q1: Array) -> Array:
+    """Multiply two quaternions (supports batching).
+
+    Args:
+        q2: Second quaternion (w, x, y, z), shape (..., 4)
+        q1: First quaternion (w, x, y, z), shape (..., 4)
+
+    Returns:
+        Product quaternion, shape (..., 4)
+    """
+    w1, x1, y1, z1 = jnp.split(q1, 4, axis=-1)
+    w2, x2, y2, z2 = jnp.split(q2, 4, axis=-1)
+
+    w = w2 * w1 - x2 * x1 - y2 * y1 - z2 * z1
+    x = w2 * x1 + x2 * w1 + y2 * z1 - z2 * y1
+    y = w2 * y1 - x2 * z1 + y2 * w1 + z2 * x1
+    z = w2 * z1 + x2 * y1 - y2 * x1 + z2 * w1
+
+    return jnp.concatenate([w, x, y, z], axis=-1)

{xax-0.2.23.dist-info → xax-0.3.1.dist-info}/METADATA

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

{xax-0.2.23.dist-info → xax-0.3.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=Q0boKxPtEUiiJ9j7Cdx51bLLFtYx3fPfCTG-o8o2Chk,15653
+xax/__init__.py,sha256=Ku51ccrDjxWuSCgmHHVnTXhloyA_d9A1VlEEbO0Ycjg,15713
 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
@@ -11,7 +11,7 @@ xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/attention.py,sha256=0essK90OO3x9FxnUqU0DhufwXKRMN41zMtRCki5iKzQ,24742
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
-xax/nn/geom.py,sha256=A7WPefMvgwUNReZC7_HX1GmvHPASyghbaXaKsuhwDrE,7382
+xax/nn/geom.py,sha256=7_cUeoJIfCS_YbQ6O7zElnpYsAn4XrVm4m7PaT9UcYw,8005
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
 xax/nn/metrics.py,sha256=zuvPXlRQczBTLHD4ilNGmZaiq6Yie3rxCMq6JkI_kos,3154
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
@@ -59,9 +59,9 @@ 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.23.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.23.dist-info/METADATA,sha256=mA98vsIjdfb8XM2mN1vUb2VRVEPU4xf10IWLxxFJjmY,1247
-xax-0.2.23.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-xax-0.2.23.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.2.23.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.23.dist-info/RECORD,,
+xax-0.3.1.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.3.1.dist-info/METADATA,sha256=xAwDUWNoga60_KZ_nPEXrdAuvlak-zWAXk3lE-XvLeI,1246
+xax-0.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xax-0.3.1.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.3.1.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.3.1.dist-info/RECORD,,

{xax-0.2.23.dist-info → xax-0.3.1.dist-info}/WHEEL

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