xax 0.3.12__tar.gz → 0.3.14__tar.gz

{xax-0.3.12/xax.egg-info → xax-0.3.14}/PKG-INFO

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

{xax-0.3.12 → xax-0.3.14}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.12"
+__version__ = "0.3.14"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -53,6 +53,7 @@ __all__ = [
     "quat_mul",
     "quat_to_euler",
     "quat_to_rotmat",
+    "quat_to_yaw",
     "rotate_vector_by_quat",
     "rotation6d_to_rotation_matrix",
     "rotation_matrix_to_quat",
@@ -100,9 +101,9 @@ __all__ = [
     "Task",
     "collate",
     "collate_non_null",
-    "breakpoint_if_nan",
+    "breakpoint_if_nonfinite",
     "get_named_leaves",
-    "log_if_nan",
+    "log_if_nonfinite",
     "BaseFileDownloader",
     "ContextTimer",
     "CumulativeTimer",
@@ -198,7 +199,10 @@ if "XLA_FLAGS" in os.environ:
 # If Nvidia GPU is detected (meaning, is `nvidia-smi` available?), disable
 # Triton GEMM kernels. See https://github.com/NVIDIA/JAX-Toolbox
 if shutil.which("nvidia-smi") is not None:
-    xla_flags += ["--xla_gpu_enable_latency_hiding_scheduler=true", "--xla_gpu_enable_triton_gemm=false"]
+    xla_flags += [
+        "--xla_gpu_enable_latency_hiding_scheduler=true",
+        "--xla_gpu_enable_triton_gemm=false",
+    ]
 os.environ["XLA_FLAGS"] = " ".join(xla_flags)
 
 # If this flag is set, eagerly imports the entire package (not recommended).
@@ -246,6 +250,7 @@ NAME_MAP: dict[str, str] = {
     "quat_mul": "nn.geom",
     "quat_to_euler": "nn.geom",
     "quat_to_rotmat": "nn.geom",
+    "quat_to_yaw": "nn.geom",
     "rotate_vector_by_quat": "nn.geom",
     "rotation6d_to_rotation_matrix": "nn.geom",
     "rotation_matrix_to_quat": "nn.geom",
@@ -293,9 +298,9 @@ NAME_MAP: dict[str, str] = {
     "Task": "task.task",
     "collate": "utils.data.collate",
     "collate_non_null": "utils.data.collate",
-    "breakpoint_if_nan": "utils.debugging",
+    "breakpoint_if_nonfinite": "utils.debugging",
     "get_named_leaves": "utils.debugging",
-    "log_if_nan": "utils.debugging",
+    "log_if_nonfinite": "utils.debugging",
     "BaseFileDownloader": "utils.experiments",
     "ContextTimer": "utils.experiments",
     "CumulativeTimer": "utils.experiments",
@@ -443,6 +448,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         quat_mul,
         quat_to_euler,
         quat_to_rotmat,
+        quat_to_yaw,
         rotate_vector_by_quat,
         rotation6d_to_rotation_matrix,
         rotation_matrix_to_quat,
@@ -486,7 +492,11 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.script import Script, ScriptConfig
     from xax.task.task import Config, Task
     from xax.utils.data.collate import CollateMode, collate, collate_non_null
-    from xax.utils.debugging import breakpoint_if_nan, get_named_leaves, log_if_nan
+    from xax.utils.debugging import (
+        breakpoint_if_nonfinite,
+        get_named_leaves,
+        log_if_nonfinite,
+    )
     from xax.utils.experiments import (
         BaseFileDownloader,
         ContextTimer,

{xax-0.3.12 → xax-0.3.14}/xax/nn/geom.py

@@ -1,6 +1,7 @@
 """Defines geometry functions."""
 
 import chex
+import jax
 from jax import numpy as jnp
 from jaxtyping import Array
 
@@ -15,30 +16,53 @@ def quat_to_euler(quat_4: Array, eps: float = 1e-6) -> Array:
     Returns:
         The roll, pitch, yaw angles with shape (*, 3).
     """
-    quat_4 = quat_4 / (jnp.linalg.norm(quat_4, axis=-1, keepdims=True) + eps)
-    w, x, y, z = jnp.split(quat_4, 4, axis=-1)
+    # Normalize with clamping
+    norm_sq = jnp.sum(quat_4**2, axis=-1, keepdims=True)
+    inv_norm = jax.lax.rsqrt(jnp.maximum(norm_sq, eps))
+    quat_4 = quat_4 * inv_norm
+
+    w, x, y, z = jnp.unstack(quat_4, axis=-1)
 
     # Roll (x-axis rotation)
     sinr_cosp = 2.0 * (w * x + y * z)
     cosr_cosp = 1.0 - 2.0 * (x * x + y * y)
-    roll = jnp.arctan2(sinr_cosp, cosr_cosp)
+    roll = jax.lax.atan2(sinr_cosp, cosr_cosp)
 
     # Pitch (y-axis rotation)
     sinp = 2.0 * (w * y - z * x)
-
-    # Handle edge cases where |sinp| >= 1
-    pitch = jnp.where(
-        jnp.abs(sinp) >= 1.0,
-        jnp.sign(sinp) * jnp.pi / 2.0,  # Use 90 degrees if out of range
-        jnp.arcsin(sinp),
-    )
+    sinp = jnp.clip(sinp, -1.0, 1.0)  # Clamp to valid domain
+    pitch = jax.lax.asin(sinp)
 
     # Yaw (z-axis rotation)
     siny_cosp = 2.0 * (w * z + x * y)
     cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
-    yaw = jnp.arctan2(siny_cosp, cosy_cosp)
+    yaw = jax.lax.atan2(siny_cosp, cosy_cosp)
+
+    return jnp.stack([roll, pitch, yaw], axis=-1)
+
+
+def quat_to_yaw(quat_4: Array, eps: float = 1e-6) -> Array:
+    """Converts a quaternion to a yaw angle.
+
+    Args:
+        quat_4: The quaternion to convert, shape (*, 4).
+        eps: A small epsilon value to avoid division by zero.
+
+    Returns:
+        The yaw angle, shape (*).
+    """
+    # Normalize using a max + safe norm to handle extremely small values robustly
+    norm_sq = jnp.sum(quat_4**2, axis=-1, keepdims=True)
+    inv_norm = jax.lax.rsqrt(jnp.maximum(norm_sq, eps))
+    quat_4 = quat_4 * inv_norm
+
+    w, x, y, z = jnp.unstack(quat_4, axis=-1)
+
+    # Compute components with clamping to avoid rounding errors near limits
+    siny_cosp = 2.0 * (w * z + x * y)
+    cosy_cosp = 1.0 - 2.0 * (y * y + z * z)
 
-    return jnp.concatenate([roll, pitch, yaw], axis=-1)
+    return jax.lax.atan2(siny_cosp, cosy_cosp)
 
 
 def euler_to_quat(euler_3: Array) -> Array:
@@ -89,7 +113,12 @@ def get_projected_gravity_vector_from_quat(quat: Array, eps: float = 1e-6) -> Ar
     return rotate_vector_by_quat(jnp.array([0, 0, -9.81]), quat, inverse=True, eps=eps)
 
 
-def rotate_vector_by_quat(vector: Array, quat: Array, inverse: bool = False, 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:

{xax-0.3.12 → xax-0.3.14}/xax/task/mixins/artifacts.py

@@ -82,7 +82,7 @@ class ArtifactsMixin(BaseTask[Config]):
             return self._exp_dir
 
         def get_exp_dir(run_id: int) -> Path:
-            return self.run_dir / f"run_{run_id}"
+            return self.run_dir / f"run_{run_id:03d}"
 
         run_id = 0
         while (exp_dir := get_exp_dir(run_id)).is_dir():

{xax-0.3.12 → xax-0.3.14}/xax/task/mixins/train.py

@@ -678,7 +678,7 @@ class TrainMixin(
 
     def log_state(self) -> None:
         logger.log(LOG_STATUS, self.task_path)
-        logger.log(LOG_STATUS, self.task_name)
+        logger.log(LOG_STATUS, self.exp_dir)
         logger.log(LOG_STATUS, "JAX devices: %s", jax.devices())
         self.logger.log_file("state.txt", get_state_file_string(self))
         self.logger.log_file("training_code.py", get_training_code(self))

{xax-0.3.12 → xax-0.3.14}/xax/utils/debugging.py

@@ -51,9 +51,25 @@ def get_named_leaves(
     return ret
 
 
-def breakpoint_if_nan(x: Array) -> None:
-    jax.lax.cond(jnp.any(jnp.isnan(x)), lambda: jax.debug.breakpoint(), lambda: None)
+def breakpoint_if_nonfinite(x: Array) -> None:
+    is_finite = jnp.isfinite(x).all()
 
+    def true_fn(x: Array) -> None:
+        pass
 
-def log_if_nan(x: Array, loc: str) -> None:
-    jax.lax.cond(jnp.any(jnp.isnan(x)), lambda: jax.debug.print("=== NaNs: {loc} ===", loc=loc), lambda: None)
+    def false_fn(x: Array) -> None:
+        jax.debug.breakpoint()
+
+    jax.lax.cond(is_finite, true_fn, false_fn, x)
+
+
+def log_if_nonfinite(x: Array, loc: str) -> None:
+    is_finite = jnp.isfinite(x).all()
+
+    def true_fn(x: Array) -> None:
+        pass
+
+    def false_fn(x: Array) -> None:
+        jax.debug.print("=== NaNs: {loc} ===", loc=loc)
+
+    jax.lax.cond(is_finite, true_fn, false_fn, x)

{xax-0.3.12 → xax-0.3.14}/xax/utils/pytree.py

@@ -274,6 +274,9 @@ def diff_pytree(tree_a: PyTree, tree_b: PyTree, prefix: str = "") -> list[str]:
 
     # Handles dataclasses.
     if is_dataclass(tree_a) and is_dataclass(tree_b):
+        if type(tree_a) is not type(tree_b):
+            diffs.append(f"{prefix}: type {type(tree_a)} vs {type(tree_b)}")
+            return diffs
         for field in fields(tree_a):
             attr_a, attr_b = getattr(tree_a, field.name), getattr(tree_b, field.name)
             diffs.extend(diff_pytree(attr_a, attr_b, prefix + f"{field.name}."))
@@ -330,10 +333,5 @@ def diff_pytree(tree_a: PyTree, tree_b: PyTree, prefix: str = "") -> list[str]:
             diffs.append(f"{prefix}: aval {aval_a} vs {aval_b}")
         return diffs
 
-    # Handle mismatched types
-    elif type(tree_a) is not type(tree_b):
-        diffs.append(f"{prefix}: type {type(tree_a)} vs {type(tree_b)}")
-        return diffs
-
     else:
         raise ValueError(f"Unknown type: {type(tree_a)}")

{xax-0.3.12 → xax-0.3.14/xax.egg-info}/PKG-INFO

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