xax 0.3.12__py3-none-any.whl → 0.3.14__py3-none-any.whl

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

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

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=HXD6tR7Bz1b5ImFyRyR1kAok-dx5g8eBDpO_lCIP8rk,16782
+xax/__init__.py,sha256=l3acv85D5Sq8IEv1tajSuCVY_eTGt8iJGnu_JuONB48,16944
 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
@@ -12,7 +12,7 @@ xax/nn/attention.py,sha256=m6yEoRqf7-wLgrEltaR6CxF_Cody0MaNtAkuKk39qJI,31176
 xax/nn/distributions.py,sha256=6YOjyiPOC7XLDaMYpFNBlLCu3eLgDAeqIg9FoKfYLL4,6497
 xax/nn/embeddings.py,sha256=8tAuAPdkVj-U5IwtRZKHA0WYMFRbpCuwyAxcChdKhbE,11784
 xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
-xax/nn/geom.py,sha256=c9K52vLm-V-15CRqMNx0OmqsWfb3PHQxXW4OSx9kCAk,10635
+xax/nn/geom.py,sha256=ataKbQFXTebK9fM10CFyxsHOPGXhn26P4jakoc9Wqek,11424
 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
@@ -33,7 +33,7 @@ xax/task/loggers/state.py,sha256=0Jy0NYnY4c0qt0LvNlaTaCKOSqk5SCKln5VdyuQGnIc,140
 xax/task/loggers/stdout.py,sha256=giKSW2R83YkgRefm3BLkE7t8Pbj5Dux4AgsdJxYIbGo,6619
 xax/task/loggers/tensorboard.py,sha256=sRyBbeBeVXDTYhPZIKIapW0JEfL9hqqzhNTeIcSd374,8883
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
-xax/task/mixins/artifacts.py,sha256=R-y3p7__zJHlHDqwDVAZysg2ZmebCJbqAx_xGT2Xpd0,3857
+xax/task/mixins/artifacts.py,sha256=UN26TW22ARduO6Bjs0yRu4-V6-Md9MPbXLKDnS28m44,3861
 xax/task/mixins/checkpointing.py,sha256=v50IZ7j58DWmEu-_6Zh_02R5KUVGhrMkg5n-MYM_J4c,11484
 xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 xax/task/mixins/cpu_stats.py,sha256=rO_9a82ZdsNec61ya4FpYE-rWqPhpijRSXsOfc6caFA,9595
@@ -43,16 +43,16 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=hqDEsMp_SL6ee97iq26-G0g49OcWZZaX82JD4F22eJU,1781
 xax/task/mixins/runnable.py,sha256=pcLrYc_TycZUY9zZim05Skc2FWk3IZKFnu6p3UDMonM,1966
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=hwAR_G1kgvhXgrE5ZRNL4Jn-Teflx65_1bdk6aULXEg,32814
+xax/task/mixins/train.py,sha256=qb0zpsyeCk_U8Sk8THxtXkUVwj5r0lOlMLNRTctvcWU,32812
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
+xax/utils/debugging.py,sha256=85JYIdnzLnvXsuli-4YHei_3tE3DnX3rmDSARKW2u1M,2192
 xax/utils/experiments.py,sha256=5k5hPYSaVjzoR_nm2Q3DAHMMYi3Bcp3N3PAQbwZq7Gg,29830
 xax/utils/jax.py,sha256=6cP95-rcjkRt1fefkZWJQhJhH0uUYWJB3w4NP1-aDp0,10136
 xax/utils/jaxpr.py,sha256=H7pWl48ROXIB1-ZPWYfOn-ou3EBMxYWIwc_A0reJQoo,2333
 xax/utils/logging.py,sha256=Kkyma_LJXqrN2HTQ214gRP_9ih3_bKk115MWC60lQWM,6656
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
-xax/utils/pytree.py,sha256=e8T5DY0ZhPcbvS3EuOsac0Oprra46lN05WEIhVN-3V0,12670
+xax/utils/pytree.py,sha256=qC7OfCydX3N5yDIgcWwiXFIdpQZg3uxgBP2H85eNmzQ,12649
 xax/utils/tensorboard.py,sha256=P0oIFvX2Qts1H4lkpizhRIpQdD0MNppVMeut0Z94yCs,19878
 xax/utils/text.py,sha256=xS02aSzdywl3KIaNSpKWcxdd37oYlUJtu9wIjkc1wVc,10654
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -60,9 +60,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.3.12.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.3.12.dist-info/METADATA,sha256=RACxHJ_iF4r0BTTTgyTI1ExYF_-aXRWrsq3NlQC7l9A,1247
-xax-0.3.12.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xax-0.3.12.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.3.12.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.3.12.dist-info/RECORD,,
+xax-0.3.14.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.3.14.dist-info/METADATA,sha256=eb-f4GhyPyCizmbj87lEg7CE7ufKtAf7uEGLI9mBms4,1247
+xax-0.3.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xax-0.3.14.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.3.14.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.3.14.dist-info/RECORD,,