xax 0.3.12__tar.gz → 0.4.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.3.12
+Version: 0.4.4
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,6 +31,8 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
+Provides-Extra: wandb
+Requires-Dist: wandb[media]; extra == "wandb"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.3.12 → xax-0.4.4}/pyproject.toml

@@ -35,6 +35,7 @@ explicit_package_bases = true
 [[tool.mypy.overrides]]
 
 module = [
+    "chex.*",
     "optax.*",
     "setuptools.*",
     "tensorboard.*",

{xax-0.3.12 → xax-0.4.4}/setup.py

@@ -33,6 +33,7 @@ setup(
     tests_require=requirements_dev,
     extras_require={
         "dev": requirements_dev,
+        "wandb": ["wandb[media]"],
     },
     package_data={
         "xax": [

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

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.3.12"
+__version__ = "0.4.4"
 
 # 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",
@@ -93,16 +94,19 @@ __all__ = [
     "DataloaderConfig",
     "GPUStatsOptions",
     "StepContext",
+    "InitParams",
     "ValidStepTimer",
     "Script",
     "ScriptConfig",
     "Config",
+    "SupervisedConfig",
+    "SupervisedTask",
     "Task",
     "collate",
     "collate_non_null",
-    "breakpoint_if_nan",
+    "breakpoint_if_nonfinite",
     "get_named_leaves",
-    "log_if_nan",
+    "log_if_nonfinite",
     "BaseFileDownloader",
     "ContextTimer",
     "CumulativeTimer",
@@ -167,6 +171,7 @@ __all__ = [
     "uncolored",
     "wrapped",
     "FrozenDict",
+    "freeze_dict",
     "HashableArray",
     "hashable_array",
 ]
@@ -198,7 +203,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 +254,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",
@@ -286,16 +295,19 @@ NAME_MAP: dict[str, str] = {
     "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
     "StepContext": "task.mixins.step_wrapper",
+    "InitParams": "task.mixins.train",
     "ValidStepTimer": "task.mixins.train",
     "Script": "task.script",
     "ScriptConfig": "task.script",
     "Config": "task.task",
+    "SupervisedConfig": "task.task",
+    "SupervisedTask": "task.task",
     "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",
@@ -360,6 +372,7 @@ NAME_MAP: dict[str, str] = {
     "uncolored": "utils.text",
     "wrapped": "utils.text",
     "FrozenDict": "utils.types.frozen_dict",
+    "freeze_dict": "utils.types.frozen_dict",
     "HashableArray": "utils.types.hashable_array",
     "hashable_array": "utils.types.hashable_array",
 }
@@ -443,6 +456,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,
@@ -482,11 +496,15 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.mixins.data_loader import DataloaderConfig
     from xax.task.mixins.gpu_stats import GPUStatsOptions
     from xax.task.mixins.step_wrapper import StepContext
-    from xax.task.mixins.train import Batch, Output, ValidStepTimer
+    from xax.task.mixins.train import Batch, InitParams, Output, ValidStepTimer
     from xax.task.script import Script, ScriptConfig
-    from xax.task.task import Config, Task
+    from xax.task.task import Config, SupervisedConfig, SupervisedTask, 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,
@@ -556,7 +574,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         uncolored,
         wrapped,
     )
-    from xax.utils.types.frozen_dict import FrozenDict
+    from xax.utils.types.frozen_dict import FrozenDict, freeze_dict
     from xax.utils.types.hashable_array import HashableArray, hashable_array
 
 del TYPE_CHECKING, IMPORT_ALL

{xax-0.3.12 → xax-0.4.4}/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.4.4/xax/task/launchers/single_process.py

@@ -0,0 +1,141 @@
+"""Defines a launcher to train a model locally, in a single process."""
+
+import logging
+import os
+import shutil
+import subprocess
+from typing import TYPE_CHECKING
+
+import jax
+
+from xax.task.base import RawConfigType
+from xax.task.launchers.base import BaseLauncher
+from xax.task.mixins.gpu_stats import get_num_gpus
+from xax.utils.logging import configure_logging
+
+if TYPE_CHECKING:
+    from xax.task.mixins.runnable import Config, RunnableMixin
+
+
+def get_gpu_memory_info() -> dict[int, tuple[float, float]]:
+    """Get memory information for all GPUs.
+
+    Returns:
+        Dictionary mapping GPU index to (total_memory_mb, used_memory_mb)
+    """
+    command = "nvidia-smi --query-gpu=index,memory.total,memory.used --format=csv,noheader"
+
+    try:
+        with subprocess.Popen(command.split(), stdout=subprocess.PIPE, universal_newlines=True) as proc:
+            stdout = proc.stdout
+            assert stdout is not None
+
+            gpu_info = {}
+            for line in stdout:
+                line = line.strip()
+                if not line:
+                    continue
+
+                parts = line.split(", ")
+                if len(parts) >= 3:
+                    gpu_id = int(parts[0])
+                    total_mem = float(parts[1].replace(" MiB", ""))
+                    used_mem = float(parts[2].replace(" MiB", ""))
+                    gpu_info[gpu_id] = (total_mem, used_mem)
+
+            return gpu_info
+
+    except Exception as e:
+        logger = configure_logging()
+        logger.warning("Failed to get GPU memory info: %s", e)
+        return {}
+
+
+def select_best_gpu() -> int | None:
+    """Select the GPU with the most available memory.
+
+    Returns:
+        GPU index with most available memory, or None if no GPUs found
+    """
+    gpu_info = get_gpu_memory_info()
+
+    if not gpu_info:
+        return None
+
+    # Find GPU with most available memory
+    best_gpu = None
+    max_available: float = -1.0
+
+    for gpu_id, (total_mem, used_mem) in gpu_info.items():
+        available_mem = total_mem - used_mem
+        if available_mem > max_available:
+            max_available = available_mem
+            best_gpu = gpu_id
+
+    return best_gpu
+
+
+def configure_gpu_devices(logger: logging.Logger | None = None) -> None:
+    if logger is None:
+        logger = configure_logging()
+
+    # If there are multiple devices, choose the one with the most
+    # available memory (i.e., the one which is likely not being used
+    # by other processes) and use only that device.
+    num_gpus = get_num_gpus()
+
+    if num_gpus > 1:
+        logger.info("Multiple GPUs detected (%d), selecting GPU with most available memory", num_gpus)
+
+        best_gpu = select_best_gpu()
+        if best_gpu is not None:
+            logger.info("Selected GPU %d for training", best_gpu)
+
+            # Set CUDA_VISIBLE_DEVICES to only show the selected GPU
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(best_gpu)
+
+            # Configure JAX to use the selected device
+            try:
+                devices = jax.devices("gpu")
+                if devices:
+                    jax.config.update("jax_default_device", devices[0])
+                    logger.info("Configured JAX to use device: %s", devices[0])
+            except Exception as e:
+                logger.warning("Failed to configure JAX device: %s", e)
+        else:
+            logger.warning("Could not determine best GPU, using default device selection")
+    elif num_gpus == 1:
+        logger.info("Single GPU detected, using default device selection")
+
+
+def configure_devices(logger: logging.Logger | None = None) -> None:
+    if logger is None:
+        logger = configure_logging()
+
+    if shutil.which("nvidia-smi") is not None:
+        configure_gpu_devices(logger)
+
+
+def run_single_process_training(
+    task: "type[RunnableMixin[Config]]",
+    *cfgs: RawConfigType,
+    use_cli: bool | list[str] = True,
+    logger: logging.Logger | None = None,
+) -> None:
+    if logger is None:
+        logger = configure_logging()
+    task_obj = task.get_task(*cfgs, use_cli=use_cli)
+    task_obj.add_logger_handlers(logger)
+    task_obj.run()
+
+
+class SingleProcessLauncher(BaseLauncher):
+    def launch(
+        self,
+        task: "type[RunnableMixin[Config]]",
+        *cfgs: RawConfigType,
+        use_cli: bool | list[str] = True,
+    ) -> None:
+        logger = configure_logging()
+        configure_devices(logger)
+        run_single_process_training(task, *cfgs, use_cli=use_cli, logger=logger)