xax 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

xax/__init__.py

@@ -11,7 +11,7 @@ This file can be maintained by running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.0.5"
+__version__ = "0.0.7"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -34,6 +34,9 @@ __all__ = [
     "get_positional_embeddings",
     "get_rotary_embeddings",
     "rotary_embeddings",
+    "euler_to_quat",
+    "quat_to_euler",
+    "is_master",
     "BaseLauncher",
     "CliLauncher",
     "SingleProcessLauncher",
@@ -56,27 +59,52 @@ __all__ = [
     "collate",
     "collate_non_null",
     "BaseFileDownloader",
+    "CumulativeTimer",
     "DataDownloader",
+    "IntervalTicker",
+    "IterationTimer",
+    "MinGradScaleError",
     "ModelDownloader",
+    "NaNError",
+    "StateTimer",
+    "TrainingFinishedError",
     "check_md5",
     "check_sha256",
+    "cpu_count",
+    "date_str",
+    "diff_configs",
     "get_git_state",
+    "get_random_port",
     "get_state_dict_prefix",
     "get_training_code",
     "save_config",
+    "stage_environment",
+    "to_markdown_table",
+    "jit",
     "ColoredFormatter",
     "configure_logging",
     "one_hot",
     "partial_flatten",
     "worker_chunk",
+    "profile",
+    "compute_nan_ratio",
+    "flatten_array",
+    "flatten_pytree",
+    "slice_array",
+    "slice_pytree",
     "TextBlock",
+    "camelcase_to_snakecase",
     "colored",
     "format_datetime",
     "format_timedelta",
+    "highlight_exception_message",
+    "is_interactive_session",
     "outlined",
     "render_text_blocks",
     "show_error",
+    "show_info",
     "show_warning",
+    "snakecase_to_camelcase",
     "uncolored",
     "wrapped",
 ]
@@ -85,8 +113,12 @@ __all__ += [
     "Batch",
     "CollateMode",
     "EmbeddingKind",
+    "LOG_ERROR_SUMMARY",
+    "LOG_PING",
+    "LOG_STATUS",
     "Output",
     "Phase",
+    "RawConfigType",
 ]
 
 import os
@@ -119,6 +151,9 @@ NAME_MAP: dict[str, str] = {
     "get_positional_embeddings": "nn.embeddings",
     "get_rotary_embeddings": "nn.embeddings",
     "rotary_embeddings": "nn.embeddings",
+    "euler_to_quat": "nn.geom",
+    "quat_to_euler": "nn.geom",
+    "is_master": "nn.parallel",
     "BaseLauncher": "task.launchers.base",
     "CliLauncher": "task.launchers.cli",
     "SingleProcessLauncher": "task.launchers.single_process",
@@ -141,27 +176,52 @@ NAME_MAP: dict[str, str] = {
     "collate": "utils.data.collate",
     "collate_non_null": "utils.data.collate",
     "BaseFileDownloader": "utils.experiments",
+    "CumulativeTimer": "utils.experiments",
     "DataDownloader": "utils.experiments",
+    "IntervalTicker": "utils.experiments",
+    "IterationTimer": "utils.experiments",
+    "MinGradScaleError": "utils.experiments",
     "ModelDownloader": "utils.experiments",
+    "NaNError": "utils.experiments",
+    "StateTimer": "utils.experiments",
+    "TrainingFinishedError": "utils.experiments",
     "check_md5": "utils.experiments",
     "check_sha256": "utils.experiments",
+    "cpu_count": "utils.experiments",
+    "date_str": "utils.experiments",
+    "diff_configs": "utils.experiments",
     "get_git_state": "utils.experiments",
+    "get_random_port": "utils.experiments",
     "get_state_dict_prefix": "utils.experiments",
     "get_training_code": "utils.experiments",
     "save_config": "utils.experiments",
+    "stage_environment": "utils.experiments",
+    "to_markdown_table": "utils.experiments",
+    "jit": "utils.jax",
     "ColoredFormatter": "utils.logging",
     "configure_logging": "utils.logging",
     "one_hot": "utils.numpy",
     "partial_flatten": "utils.numpy",
     "worker_chunk": "utils.numpy",
+    "profile": "utils.profile",
+    "compute_nan_ratio": "utils.pytree",
+    "flatten_array": "utils.pytree",
+    "flatten_pytree": "utils.pytree",
+    "slice_array": "utils.pytree",
+    "slice_pytree": "utils.pytree",
     "TextBlock": "utils.text",
+    "camelcase_to_snakecase": "utils.text",
     "colored": "utils.text",
     "format_datetime": "utils.text",
     "format_timedelta": "utils.text",
+    "highlight_exception_message": "utils.text",
+    "is_interactive_session": "utils.text",
     "outlined": "utils.text",
     "render_text_blocks": "utils.text",
     "show_error": "utils.text",
+    "show_info": "utils.text",
     "show_warning": "utils.text",
+    "snakecase_to_camelcase": "utils.text",
     "uncolored": "utils.text",
     "wrapped": "utils.text",
 }
@@ -172,8 +232,12 @@ NAME_MAP.update(
         "Batch": "task.mixins.train",
         "CollateMode": "utils.data.collate",
         "EmbeddingKind": "nn.embeddings",
+        "LOG_ERROR_SUMMARY": "utils.logging",
+        "LOG_PING": "utils.logging",
+        "LOG_STATUS": "utils.logging",
         "Output": "task.mixins.output",
         "Phase": "core.state",
+        "RawConfigType": "task.base",
     },
 )
 
@@ -211,6 +275,9 @@ if IMPORT_ALL or TYPE_CHECKING:
         get_rotary_embeddings,
         rotary_embeddings,
     )
+    from xax.nn.geom import euler_to_quat, quat_to_euler
+    from xax.nn.parallel import is_master
+    from xax.task.base import RawConfigType
     from xax.task.launchers.base import BaseLauncher
     from xax.task.launchers.cli import CliLauncher
     from xax.task.launchers.single_process import SingleProcessLauncher
@@ -229,26 +296,59 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.utils.data.collate import CollateMode, collate, collate_non_null
     from xax.utils.experiments import (
         BaseFileDownloader,
+        CumulativeTimer,
         DataDownloader,
+        IntervalTicker,
+        IterationTimer,
+        MinGradScaleError,
         ModelDownloader,
+        NaNError,
+        StateTimer,
+        TrainingFinishedError,
         check_md5,
         check_sha256,
+        cpu_count,
+        date_str,
+        diff_configs,
         get_git_state,
+        get_random_port,
         get_state_dict_prefix,
         get_training_code,
         save_config,
+        stage_environment,
+        to_markdown_table,
+    )
+    from xax.utils.jax import jit
+    from xax.utils.logging import (
+        LOG_ERROR_SUMMARY,
+        LOG_PING,
+        LOG_STATUS,
+        ColoredFormatter,
+        configure_logging,
     )
-    from xax.utils.logging import ColoredFormatter, configure_logging
     from xax.utils.numpy import one_hot, partial_flatten, worker_chunk
+    from xax.utils.profile import profile
+    from xax.utils.pytree import (
+        compute_nan_ratio,
+        flatten_array,
+        flatten_pytree,
+        slice_array,
+        slice_pytree,
+    )
     from xax.utils.text import (
         TextBlock,
+        camelcase_to_snakecase,
         colored,
         format_datetime,
         format_timedelta,
+        highlight_exception_message,
+        is_interactive_session,
         outlined,
         render_text_blocks,
         show_error,
+        show_info,
         show_warning,
+        snakecase_to_camelcase,
         uncolored,
         wrapped,
     )

xax/core/conf.py

@@ -6,7 +6,6 @@ from dataclasses import dataclass, field as field_base
 from pathlib import Path
 from typing import Any, cast
 
-import jax.numpy as jnp
 from omegaconf import II, MISSING, Container as OmegaConfContainer, OmegaConf
 
 from xax.utils.text import show_error
@@ -61,68 +60,44 @@ def is_missing(cfg: Any, key: str) -> bool:  # noqa: ANN401
     return False
 
 
-@dataclass
+@dataclass(kw_only=True)
 class Logging:
     hide_third_party_logs: bool = field(True, help="If set, hide third-party logs")
-    log_level: str = field("INFO", help="The logging level to use")
+    log_level: str = field(II("oc.env:XAX_LOG_LEVEL,INFO"), help="The logging level to use")
 
 
-@dataclass
-class Device:
-    cpu: bool = field(True, help="Whether to use the CPU")
-    gpu: bool = field(II("oc.env:USE_GPU,1"), help="Whether to use the GPU")
-    metal: bool = field(II("oc.env:USE_METAL,1"), help="Whether to use the Apple Silicon accelerator")
-    use_fp64: bool = field(False, help="Always use the 64-bit floating point type")
-    use_fp32: bool = field(False, help="Always use the 32-bit floating point type")
-    use_bf16: bool = field(False, help="Always use the 16-bit bfloat type")
-    use_fp16: bool = field(False, help="Always use the 16-bit floating point type")
-
-
-def parse_dtype(cfg: Device) -> jnp.dtype | None:
-    if cfg.use_fp64:
-        return jnp.float64
-    if cfg.use_fp32:
-        return jnp.float32
-    if cfg.use_bf16:
-        return jnp.bfloat16
-    if cfg.use_fp16:
-        return jnp.float16
-    return None
-
-
-@dataclass
+@dataclass(kw_only=True)
 class Triton:
     use_triton_if_available: bool = field(True, help="Use Triton if available")
 
 
-@dataclass
+@dataclass(kw_only=True)
 class Experiment:
     default_random_seed: int = field(1337, help="The default random seed to use")
     max_workers: int = field(32, help="Maximum number of workers to use")
 
 
-@dataclass
+@dataclass(kw_only=True)
 class Directories:
     run: str = field(II("oc.env:RUN_DIR"), help="The run directory")
     data: str = field(II("oc.env:DATA_DIR"), help="The data directory")
     pretrained_models: str = field(II("oc.env:MODEL_DIR"), help="The models directory")
 
 
-@dataclass
+@dataclass(kw_only=True)
 class SlurmPartition:
     partition: str = field(MISSING, help="The partition name")
     num_nodes: int = field(1, help="The number of nodes to use")
 
 
-@dataclass
+@dataclass(kw_only=True)
 class Slurm:
     launch: dict[str, SlurmPartition] = field({}, help="The available launch configurations")
 
 
-@dataclass
+@dataclass(kw_only=True)
 class UserConfig:
     logging: Logging = field(Logging)
-    device: Device = field(Device)
     triton: Triton = field(Triton)
     experiment: Experiment = field(Experiment)
     directories: Directories = field(Directories)

xax/core/state.py

@@ -2,7 +2,7 @@
 
 import time
 from dataclasses import dataclass
-from typing import Literal, TypedDict, cast, get_args
+from typing import Literal, NotRequired, TypedDict, cast, get_args
 
 from omegaconf import MISSING
 
@@ -18,16 +18,16 @@ def cast_phase(raw_phase: str) -> Phase:
 
 
 class StateDict(TypedDict, total=False):
-    num_steps: int
-    num_samples: int
-    num_valid_steps: int
-    num_valid_samples: int
-    start_time_s: float
-    elapsed_time_s: float
-    raw_phase: str
+    num_steps: NotRequired[int]
+    num_samples: NotRequired[int]
+    num_valid_steps: NotRequired[int]
+    num_valid_samples: NotRequired[int]
+    start_time_s: NotRequired[float]
+    elapsed_time_s: NotRequired[float]
+    raw_phase: NotRequired[str]
 
 
-@dataclass(frozen=True)
+@dataclass
 class State:
     num_steps: int = field(MISSING, help="Number of steps so far")
     num_samples: int = field(MISSING, help="Number of sample so far")
@@ -41,6 +41,10 @@ class State:
     def phase(self) -> Phase:
         return cast_phase(self.raw_phase)
 
+    @phase.setter
+    def phase(self, phase: Phase) -> None:
+        self.raw_phase = phase
+
     @classmethod
     def init_state(cls) -> "State":
         return cls(
@@ -65,17 +69,3 @@ class State:
                 return self.num_valid_steps
             case _:
                 raise ValueError(f"Invalid phase: {phase}")
-
-    def replace(self, values: StateDict) -> "State":
-        return State(
-            num_steps=values.get("num_steps", self.num_steps),
-            num_samples=values.get("num_samples", self.num_samples),
-            num_valid_steps=values.get("num_valid_steps", self.num_valid_steps),
-            num_valid_samples=values.get("num_valid_samples", self.num_valid_samples),
-            start_time_s=values.get("start_time_s", self.start_time_s),
-            elapsed_time_s=values.get("elapsed_time_s", self.elapsed_time_s),
-            raw_phase=values.get("raw_phase", self.raw_phase),
-        )
-
-    def with_phase(self, phase: Phase) -> "State":
-        return self.replace({"raw_phase": phase})

xax/nn/geom.py

@@ -0,0 +1,75 @@
+"""Defines geometry functions."""
+
+import jax
+from jax import numpy as jnp
+
+
+def quat_to_euler(quat_4: jax.Array, eps: float = 1e-6) -> jax.Array:
+    """Normalizes and converts a quaternion (w, x, y, z) to roll, pitch, yaw.
+
+    Args:
+        quat_4: The quaternion to convert, shape (*, 4).
+        eps: A small epsilon value to avoid division by zero.
+
+    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)
+
+    # 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)
+
+    # 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),
+    )
+
+    # 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)
+
+    return jnp.concatenate([roll, pitch, yaw], axis=-1)
+
+
+def euler_to_quat(euler_3: jax.Array) -> jax.Array:
+    """Converts roll, pitch, yaw angles to a quaternion (w, x, y, z).
+
+    Args:
+        euler_3: The roll, pitch, yaw angles, shape (*, 3).
+
+    Returns:
+        The quaternion with shape (*, 4).
+    """
+    # Extract roll, pitch, yaw from input
+    roll, pitch, yaw = jnp.split(euler_3, 3, axis=-1)
+
+    # Calculate trigonometric functions for each angle
+    cr = jnp.cos(roll * 0.5)
+    sr = jnp.sin(roll * 0.5)
+    cp = jnp.cos(pitch * 0.5)
+    sp = jnp.sin(pitch * 0.5)
+    cy = jnp.cos(yaw * 0.5)
+    sy = jnp.sin(yaw * 0.5)
+
+    # Calculate quaternion components using the conversion formula
+    w = cr * cp * cy + sr * sp * sy
+    x = sr * cp * cy - cr * sp * sy
+    y = cr * sp * cy + sr * cp * sy
+    z = cr * cp * sy - sr * sp * cy
+
+    # Combine into quaternion [w, x, y, z]
+    quat = jnp.concatenate([w, x, y, z], axis=-1)
+
+    # Normalize the quaternion
+    quat = quat / jnp.linalg.norm(quat, axis=-1, keepdims=True)
+
+    return quat

xax/requirements.txt

@@ -6,6 +6,8 @@ jaxtyping
 equinox
 optax
 dpshdl
+chex
+importlib-resources
 
 # Data processing and serialization
 cloudpickle

xax/task/base.py

@@ -15,6 +15,7 @@ from pathlib import Path
 from types import TracebackType
 from typing import Generic, Self, TypeVar, cast
 
+import jax
 from omegaconf import Container, DictConfig, OmegaConf
 
 from xax.core.state import State
@@ -23,6 +24,7 @@ from xax.utils.text import camelcase_to_snakecase
 logger = logging.getLogger(__name__)
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class BaseConfig:
     pass