xax 0.3.2__tar.gz → 0.3.3__tar.gz

{xax-0.3.2/xax.egg-info → xax-0.3.3}/PKG-INFO

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

{xax-0.3.2 → xax-0.3.3}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.2"
+__version__ = "0.3.3"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -42,11 +42,12 @@ __all__ = [
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "normalize",
+    "quat_mul",
     "quat_to_euler",
     "quat_to_rotmat",
-    "quat_mul",
     "rotate_vector_by_quat",
     "rotation6d_to_rotation_matrix",
+    "rotation_matrix_to_quat",
     "rotation_matrix_to_rotation6d",
     "cross_entropy",
     "cast_norm_type",
@@ -224,11 +225,12 @@ NAME_MAP: dict[str, str] = {
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "normalize": "nn.geom",
+    "quat_mul": "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_quat": "nn.geom",
     "rotation_matrix_to_rotation6d": "nn.geom",
     "cross_entropy": "nn.losses",
     "cast_norm_type": "nn.metrics",
@@ -405,6 +407,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         quat_to_rotmat,
         rotate_vector_by_quat,
         rotation6d_to_rotation_matrix,
+        rotation_matrix_to_quat,
         rotation_matrix_to_rotation6d,
     )
     from xax.nn.losses import cross_entropy

{xax-0.3.2 → xax-0.3.3}/xax/nn/geom.py

@@ -272,3 +272,85 @@ def quat_mul(q2: Array, q1: Array) -> Array:
     z = w2 * z1 + x2 * y1 - y2 * x1 + z2 * w1
 
     return jnp.concatenate([w, x, y, z], axis=-1)
+
+
+def rotation_matrix_to_quat(rotation_matrix: Array, eps: float = 1e-6) -> Array:
+    """Converts a rotation matrix to a unit quaternion ``(w, x, y, z)``.
+
+    Args:
+        rotation_matrix: The rotation matrix, shape ``(*, 3, 3)``.
+        eps: A small epsilon value to avoid division by zero when normalising.
+
+    Returns:
+        A quaternion with shape ``(*, 4)``.
+    """
+    chex.assert_shape(rotation_matrix, (..., 3, 3))
+
+    m00 = rotation_matrix[..., 0, 0]
+    m01 = rotation_matrix[..., 0, 1]
+    m02 = rotation_matrix[..., 0, 2]
+    m10 = rotation_matrix[..., 1, 0]
+    m11 = rotation_matrix[..., 1, 1]
+    m12 = rotation_matrix[..., 1, 2]
+    m20 = rotation_matrix[..., 2, 0]
+    m21 = rotation_matrix[..., 2, 1]
+    m22 = rotation_matrix[..., 2, 2]
+
+    trace = m00 + m11 + m22
+
+    # Case 0: trace is positive
+    s0 = jnp.sqrt(jnp.clip(trace + 1.0, a_min=0.0)) * 2.0  # S = 4 * qw
+    w0 = 0.25 * s0
+    x0 = (m21 - m12) / jnp.where(s0 < eps, 1.0, s0)
+    y0 = (m02 - m20) / jnp.where(s0 < eps, 1.0, s0)
+    z0 = (m10 - m01) / jnp.where(s0 < eps, 1.0, s0)
+
+    # Case 1: m00 is the largest diagonal term
+    s1 = jnp.sqrt(jnp.clip(1.0 + m00 - m11 - m22, a_min=0.0)) * 2.0  # S = 4 * qx
+    w1 = (m21 - m12) / jnp.where(s1 < eps, 1.0, s1)
+    x1 = 0.25 * s1
+    y1 = (m01 + m10) / jnp.where(s1 < eps, 1.0, s1)
+    z1 = (m02 + m20) / jnp.where(s1 < eps, 1.0, s1)
+
+    # Case 2: m11 is the largest diagonal term
+    s2 = jnp.sqrt(jnp.clip(1.0 + m11 - m00 - m22, a_min=0.0)) * 2.0  # S = 4 * qy
+    w2 = (m02 - m20) / jnp.where(s2 < eps, 1.0, s2)
+    x2 = (m01 + m10) / jnp.where(s2 < eps, 1.0, s2)
+    y2 = 0.25 * s2
+    z2 = (m12 + m21) / jnp.where(s2 < eps, 1.0, s2)
+
+    # Case 3: m22 is the largest diagonal term
+    s3 = jnp.sqrt(jnp.clip(1.0 + m22 - m00 - m11, a_min=0.0)) * 2.0  # S = 4 * qz
+    w3 = (m10 - m01) / jnp.where(s3 < eps, 1.0, s3)
+    x3 = (m02 + m20) / jnp.where(s3 < eps, 1.0, s3)
+    y3 = (m12 + m21) / jnp.where(s3 < eps, 1.0, s3)
+    z3 = 0.25 * s3
+
+    cond0 = trace > 0.0
+    cond1 = (m00 > m11) & (m00 > m22)
+    cond2 = m11 > m22
+
+    w = jnp.where(
+        cond0,
+        w0,
+        jnp.where(cond1, w1, jnp.where(cond2, w2, w3)),
+    )
+    x = jnp.where(
+        cond0,
+        x0,
+        jnp.where(cond1, x1, jnp.where(cond2, x2, x3)),
+    )
+    y = jnp.where(
+        cond0,
+        y0,
+        jnp.where(cond1, y1, jnp.where(cond2, y2, y3)),
+    )
+    z = jnp.where(
+        cond0,
+        z0,
+        jnp.where(cond1, z1, jnp.where(cond2, z2, z3)),
+    )
+
+    quat = jnp.stack([w, x, y, z], axis=-1)
+    quat = quat / (jnp.linalg.norm(quat, axis=-1, keepdims=True) + eps)
+    return quat

{xax-0.3.2 → xax-0.3.3}/xax/task/mixins/runnable.py

@@ -10,6 +10,7 @@ import jax
 
 from xax.task.base import BaseConfig, BaseTask, RawConfigType
 from xax.task.launchers.base import BaseLauncher
+from xax.task.launchers.cli import CliLauncher
 
 
 @jax.tree_util.register_dataclass
@@ -45,8 +46,6 @@ class RunnableMixin(BaseTask[Config], ABC):
         use_cli: bool | list[str] = True,
     ) -> None:
         if launcher is None:
-            from xax.task.launchers.cli import CliLauncher
-
             launcher = CliLauncher()
         launcher.launch(cls, *cfgs, use_cli=use_cli)
 

{xax-0.3.2 → xax-0.3.3/xax.egg-info}/PKG-INFO

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