xtrax 0.2.0__py3-none-any.whl

xtrax/__init__.py

@@ -0,0 +1,120 @@
+__version__ = "0.2.0"
+
+__all__ = [
+    # Core training
+    "Trainer",
+    "SafetyTrainStep",
+    "create_train_step",
+    "LossFunction",
+    "Callback",
+    "ResumableState",
+    "WeightedLoss",
+    "MultiTaskLoss",
+    "make_optimizer",
+    "adamw_with_schedule",
+    # Engine and IO
+    "Engine",
+    "BoundedCallbackHandler",
+    "save_checkpoint",
+    "load_checkpoint",
+    # Data
+    "DataModule",
+    "create_distributed_pipeline",
+    # Tiling
+    "AxisSpec",
+    "AxisDecision",
+    "BatchPlan",
+    "BatchPlanner",
+    "Vmap",
+    "SafeMap",
+    "DedupGather",
+    # Sparse
+    "SparseConfig",
+    "SparsePolicy",
+    "SparseMaskManager",
+    "sparsify_model",
+    "make_sparse_forward_fn",
+    "sparse_filter_jit",
+    # Distributed
+    "init_dist",
+    "is_distributed",
+    "LogicalMesh",
+    "with_manual_axes",
+    # Transforms
+    "safe_map",
+    "safe_scan",
+    # Safety
+    "safe_norm",
+    "safe_reciprocal",
+    # Stages
+    "TransformFn",
+    "RollingFn",
+    "FuseFn",
+]
+
+_LAZY = {
+    # Training subpackage
+    "Trainer": "xtrax.training",
+    "SafetyTrainStep": "xtrax.training",
+    "create_train_step": "xtrax.training",
+    "LossFunction": "xtrax.training",
+    "Callback": "xtrax.training",
+    "ResumableState": "xtrax.training",
+    "WeightedLoss": "xtrax.training",
+    "MultiTaskLoss": "xtrax.training",
+    "make_optimizer": "xtrax.training",
+    "adamw_with_schedule": "xtrax.training",
+    # Engine subpackage
+    "Engine": "xtrax.engine",
+    "BoundedCallbackHandler": "xtrax.engine",
+    # Checkpoint subpackage
+    "save_checkpoint": "xtrax.checkpoint",
+    "load_checkpoint": "xtrax.checkpoint",
+    # Data subpackage
+    "DataModule": "xtrax.data",
+    "create_distributed_pipeline": "xtrax.data",
+    # Tiling subpackage
+    "AxisSpec": "xtrax.tiling",
+    "AxisDecision": "xtrax.tiling",
+    "BatchPlan": "xtrax.tiling",
+    "BatchPlanner": "xtrax.tiling",
+    "Vmap": "xtrax.tiling",
+    "SafeMap": "xtrax.tiling",
+    "DedupGather": "xtrax.tiling",
+    # Sparse subpackage
+    "SparseConfig": "xtrax.sparse",
+    "SparsePolicy": "xtrax.sparse",
+    "SparseMaskManager": "xtrax.sparse",
+    "sparsify_model": "xtrax.sparse",
+    "make_sparse_forward_fn": "xtrax.sparse",
+    "sparse_filter_jit": "xtrax.sparse",
+    # Distributed subpackage
+    "init_dist": "xtrax.distributed",
+    "is_distributed": "xtrax.distributed",
+    "LogicalMesh": "xtrax.distributed",
+    "with_manual_axes": "xtrax.distributed",
+    # Transforms subpackage
+    "safe_map": "xtrax.transforms",
+    "safe_scan": "xtrax.transforms",
+    # Safety subpackage
+    "safe_norm": "xtrax.safety",
+    "safe_reciprocal": "xtrax.safety",
+    # Stages subpackage
+    "TransformFn": "xtrax.stages",
+    "RollingFn": "xtrax.stages",
+    "FuseFn": "xtrax.stages",
+}
+
+
+def __getattr__(name):
+    """Lazy import on attribute access."""
+    if name in _LAZY:
+        import importlib
+
+        return getattr(importlib.import_module(_LAZY[name]), name)
+    raise AttributeError(f"module 'xtrax' has no attribute {name!r}")
+
+
+def __dir__():
+    """Return sorted list of public attributes."""
+    return sorted(__all__)

xtrax/checkpoint/__init__.py

@@ -0,0 +1,13 @@
+"""Checkpoint utilities for xtrax."""
+
+from xtrax.checkpoint.orbax import (
+    get_checkpoint_manager,
+    load_checkpoint,
+    save_checkpoint,
+)
+
+__all__ = [
+    "get_checkpoint_manager",
+    "save_checkpoint",
+    "load_checkpoint",
+]

xtrax/checkpoint/orbax.py

@@ -0,0 +1,100 @@
+"""Checkpoint utilities using Orbax."""
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import orbax.checkpoint as ocp
+
+if TYPE_CHECKING:
+    from xtrax.training.types import ResumableState
+
+
+def get_checkpoint_manager(
+    directory: str | Path,
+    max_to_keep: int = 5,
+    keep_period: int | None = None,
+) -> ocp.CheckpointManager:
+    """Create and return a CheckpointManager for the given directory.
+
+    Args:
+        directory: Directory to store checkpoints.
+        max_to_keep: Maximum number of checkpoints to keep. Defaults to 5.
+        keep_period: Period (in steps) for keeping checkpoints. Defaults to None.
+
+    Returns:
+        A CheckpointManager instance with PyTreeCheckpointHandler configured.
+    """
+    directory = Path(directory)
+
+    # Create options for the checkpoint manager
+    options = ocp.CheckpointManagerOptions(
+        max_to_keep=max_to_keep,
+        keep_period=keep_period,
+    )
+
+    # Create the manager with PyTreeCheckpointHandler
+    manager = ocp.CheckpointManager(
+        directory,
+        options=options,
+        item_handlers={"state": ocp.PyTreeCheckpointHandler()},
+    )
+
+    return manager
+
+
+def save_checkpoint(
+    manager: ocp.CheckpointManager,
+    state: "ResumableState",
+    step: int | None = None,
+) -> None:
+    """Save a checkpoint to the manager.
+
+    Args:
+        manager: CheckpointManager instance.
+        state: ResumableState to save.
+        step: Step number to save at. If None, uses int(state.step).
+
+    Note:
+        This function must be called outside JAX-traced contexts.
+        Calls manager.wait_until_finished() after saving.
+    """
+    # Extract step: use provided step or convert state.step to int
+    step_int = step if step is not None else int(state.step)
+
+    # Save the checkpoint
+    manager.save(step=step_int, items={"state": state})
+
+    # Wait for save to complete
+    manager.wait_until_finished()
+
+
+def load_checkpoint(
+    manager: ocp.CheckpointManager,
+    state_template: "ResumableState",
+    step: int | None = None,
+) -> "ResumableState":
+    """Load a checkpoint from the manager.
+
+    Args:
+        manager: CheckpointManager instance.
+        state_template: Template ResumableState for pytree structure.
+            Required for orbax to reconstruct the pytree.
+        step: Step number to load. If None, uses latest_step().
+
+    Returns:
+        The loaded ResumableState.
+
+    Raises:
+        FileNotFoundError: If checkpoint doesn't exist or directory is empty.
+    """
+    # Determine which step to load
+    if step is None:
+        step = manager.latest_step()
+        if step is None:
+            raise FileNotFoundError("No checkpoints found in directory")
+
+    # Restore the checkpoint
+    loaded = manager.restore(step=step, items={"state": state_template})
+
+    # Extract the state from the composite result
+    return loaded["state"]

xtrax/data/__init__.py

@@ -0,0 +1,9 @@
+"""Data module for dataset loading and preprocessing."""
+
+from xtrax.data.module import DataModule
+from xtrax.data.pipeline import create_distributed_pipeline
+
+__all__ = [
+    "DataModule",
+    "create_distributed_pipeline",
+]

xtrax/data/module.py

@@ -0,0 +1,37 @@
+from collections.abc import Callable, Iterator
+from typing import Any
+
+import equinox as eqx
+
+# Module-level flag for distributed init state (real integration with
+# init_dist deferred to Phase 5/6)
+_dist_initialized: bool = False
+
+
+def _mark_dist_initialized() -> None:
+    """Call after init_dist() to allow DataModule iterators to proceed."""
+    global _dist_initialized
+    _dist_initialized = True
+
+
+class DataModule(eqx.Module):
+    dataset: Any
+    batch_size: int = eqx.field(static=True)
+    num_epochs: int | None = eqx.field(static=True)  # None = cycle indefinitely
+    seed: int = eqx.field(static=True)
+    distributed: bool = eqx.field(static=True)
+    collate_fn: Callable | None = eqx.field(static=True, default=None)
+
+    def train_iter(self) -> Iterator[Any]:
+        if self.distributed and not _dist_initialized:
+            raise RuntimeError(
+                "DataModule: distributed=True requires init_dist() before train_iter()."
+            )
+        yield from self.dataset
+
+    def eval_iter(self) -> Iterator[Any]:
+        if self.distributed and not _dist_initialized:
+            raise RuntimeError(
+                "DataModule: distributed=True requires init_dist() before eval_iter()."
+            )
+        yield from self.dataset

xtrax/data/pipeline.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+
+def create_distributed_pipeline(
+    dataset: Any,
+    global_batch_size: int,
+    num_devices: int,
+    seed: int,
+) -> Any:
+    """Create distributed pipeline with per-device batch size validation.
+
+    Raises ValueError if global_batch_size is not divisible by num_devices.
+    Stub implementation — real grain sharding deferred to Phase 5/6.
+    """
+    if global_batch_size % num_devices != 0:
+        raise ValueError(
+            f"create_distributed_pipeline: global_batch_size={global_batch_size} "
+            f"must be divisible by num_devices={num_devices}."
+        )
+    return dataset

xtrax/distributed/__init__.py

@@ -0,0 +1,16 @@
+"""Distributed training utilities for JAX."""
+
+from xtrax.distributed.init import init_dist, is_distributed
+from xtrax.distributed.sharding import (
+    ShardingPolicy,
+    get_device_mesh,
+    get_hardware_mesh_profile,
+)
+
+__all__ = [
+    "init_dist",
+    "is_distributed",
+    "ShardingPolicy",
+    "get_device_mesh",
+    "get_hardware_mesh_profile",
+]

xtrax/distributed/init.py

@@ -0,0 +1,153 @@
+"""Distributed initialization for JAX multi-process training."""
+
+import os
+from typing import Any
+
+import jax.distributed
+
+# Module-level state: stores initialization args
+_init_state: dict[str, Any] = {
+    "initialized": False,
+    "coordinator_address": None,
+    "num_processes": None,
+    "process_id": None,
+}
+
+
+def _derive_coordinator_from_nodelist(nodelist: str) -> str:
+    """Extract first node from SLURM_JOB_NODELIST and return as coordinator address."""
+    # nodelist can be "node-001,node-002,node-003" or "node-001" or "node-[001-003]"
+    # For simplicity, extract the first component (comma-separated or bracket-expanded)
+    first_node = nodelist.split(",")[0]
+    # Remove bracket ranges if present (e.g., "node-[001-003]" -> "node-[001-003]")
+    # For now, just use the first node as-is; return with default JAX port
+    return f"{first_node}:1234"
+
+
+def init_dist(
+    *,
+    coordinator_address: str | None = None,
+    num_processes: int | None = None,
+    process_id: int | None = None,
+) -> None:
+    """
+    Initialize JAX distributed training with optional auto-discovery.
+
+    Args:
+        coordinator_address: Address of the coordinator (e.g., "localhost:1234").
+                            If None, attempts SLURM env discovery or falls back
+                            to localhost.
+        num_processes: Total number of processes in the distributed setup.
+                      If None, attempts SLURM env discovery or falls back to 1.
+        process_id: This process's ID (0-indexed).
+                   If None, attempts SLURM env discovery or falls back to 0.
+
+    Raises:
+        RuntimeError: If init_dist has already been called with different arguments.
+
+    Notes:
+        - Idempotent: same args on repeat call → no-op.
+        - Different args on second call → RuntimeError.
+        - SLURM discovery: uses SLURM_PROCID, SLURM_NTASKS, SLURM_JOB_NODELIST.
+        - Single-process fallback: skips jax.distributed.initialize for
+          num_processes==1 (spec deviation: see below).
+        - After successful init, calls _mark_dist_initialized() to enable
+          DataModule(distributed=True).
+
+    Spec Deviation (§3.22):
+        The spec states that localhost fallback should call jax.distributed.initialize
+        with single-process config. However, calling jax.distributed.initialize with
+        num_processes=1 can be unstable in some environments. This implementation skips
+        the call for num_processes==1, instead calling _mark_dist_initialized() and
+        returning. If future requirements demand the single-process initialize call,
+        this can be made a parameter.
+    """
+    global _init_state
+
+    # Auto-discovery: explicit args → SLURM env → localhost fallback
+    if coordinator_address is None or num_processes is None or process_id is None:
+        # Try SLURM environment variables
+        slurm_ntasks = os.environ.get("SLURM_NTASKS")
+        slurm_procid = os.environ.get("SLURM_PROCID")
+        slurm_nodelist = os.environ.get("SLURM_JOB_NODELIST")
+
+        if slurm_ntasks is not None and slurm_procid is not None:
+            # SLURM env found: use it to fill in missing args
+            if num_processes is None:
+                num_processes = int(slurm_ntasks)
+            if process_id is None:
+                process_id = int(slurm_procid)
+            if coordinator_address is None:
+                if slurm_nodelist:
+                    coordinator_address = _derive_coordinator_from_nodelist(
+                        slurm_nodelist
+                    )
+                else:
+                    coordinator_address = "localhost:1234"
+        else:
+            # No SLURM env: use localhost fallback
+            if coordinator_address is None:
+                coordinator_address = "localhost:1234"
+            if num_processes is None:
+                num_processes = 1
+            if process_id is None:
+                process_id = 0
+
+    # Check idempotency: if already initialized, verify args match
+    if _init_state["initialized"]:
+        if (
+            _init_state["coordinator_address"] != coordinator_address
+            or _init_state["num_processes"] != num_processes
+            or _init_state["process_id"] != process_id
+        ):
+            raise RuntimeError(
+                f"init_dist already initialized with different args. "
+                f"Previous: coordinator_address="
+                f"{_init_state['coordinator_address']}, "
+                f"num_processes={_init_state['num_processes']}, "
+                f"process_id={_init_state['process_id']}. "
+                f"New: coordinator_address={coordinator_address}, "
+                f"num_processes={num_processes}, process_id={process_id}."
+            )
+        # Same args: idempotent, no-op
+        return
+
+    # Store state
+    _init_state["coordinator_address"] = coordinator_address
+    _init_state["num_processes"] = num_processes
+    _init_state["process_id"] = process_id
+
+    # Initialize JAX distributed if num_processes > 1
+    if num_processes > 1:
+        jax.distributed.initialize(
+            coordinator_address=coordinator_address,
+            num_processes=num_processes,
+            process_id=process_id,
+        )
+
+    # Mark distributed initialization complete (enables DataModule(distributed=True))
+    from xtrax.data.module import _mark_dist_initialized
+
+    _mark_dist_initialized()
+
+    _init_state["initialized"] = True
+
+
+def is_distributed() -> bool:
+    """Return True if init_dist has been called, False otherwise."""
+    return _init_state["initialized"]
+
+
+def _reset_for_testing() -> None:
+    """Reset internal state for testing. Not for production use."""
+    global _init_state
+    _init_state = {
+        "initialized": False,
+        "coordinator_address": None,
+        "num_processes": None,
+        "process_id": None,
+    }
+    # Also reset DataModule's distributed flag
+    import xtrax.data.module as dm
+
+    dm._dist_initialized = False

xtrax/distributed/sharding.py

@@ -0,0 +1,174 @@
+"""Distributed sharding utilities for JAX arrays and pytrees."""
+
+import math
+import re
+from typing import Any
+
+import equinox as eqx
+import jax
+import jax.sharding
+import jax.tree_util
+
+
+class ShardingPolicy(eqx.Module):
+    """Policy for assigning PartitionSpec to pytree leaves.
+
+    Uses regex pattern matching where first match wins. Matching is done via re.search,
+    so patterns can match anywhere in the path.
+
+    Attributes:
+        rules: Ordered tuple of (regex_pattern, PartitionSpec) pairs.
+    """
+
+    rules: tuple[tuple[str, jax.sharding.PartitionSpec], ...] = eqx.field(static=True)
+
+    def get_partition_spec(self, path: str) -> jax.sharding.PartitionSpec:
+        """Get the PartitionSpec for a given path by matching against rules in order.
+
+        Args:
+            path: String path to a pytree leaf (e.g., "layer_0/weight" or "bias").
+
+        Returns:
+            The PartitionSpec for the first matching rule, or PartitionSpec() if no
+            rule matches (defaults to fully replicated).
+        """
+        for pattern, partition_spec in self.rules:
+            if re.search(pattern, path):
+                return partition_spec
+        return jax.sharding.PartitionSpec()
+
+    def apply_to_pytree(self, pytree: Any) -> Any:
+        """Apply sharding policy to all leaves of a pytree.
+
+        Args:
+            pytree: A pytree (dict, list, etc.) with any leaf values.
+
+        Returns:
+            A pytree with the same structure but with leaves replaced by PartitionSpec
+            objects derived from their paths.
+        """
+
+        def apply_to_leaf(path, leaf):
+            # Convert jax.tree_util path to a string representation
+            path_str = self._path_to_string(path)
+            return self.get_partition_spec(path_str)
+
+        # Use jax.tree_util.tree_map_with_path to traverse with path tracking
+        # This preserves the pytree structure and applies the function to each leaf
+        return jax.tree_util.tree_map_with_path(apply_to_leaf, pytree)
+
+    @staticmethod
+    def _path_to_string(path: tuple) -> str:
+        """Convert a jax.tree_util path to a string representation.
+
+        Args:
+            path: A tuple of GetAttrKey, DictKey, SequenceKey, etc.
+
+        Returns:
+            A string representation of the path (e.g., "layer/weight").
+        """
+        parts = []
+        for key in path:
+            if isinstance(key, jax.tree_util.DictKey):
+                parts.append(str(key.key))
+            elif isinstance(key, jax.tree_util.GetAttrKey):
+                parts.append(key.name)
+            elif isinstance(key, jax.tree_util.SequenceKey):
+                parts.append(f"[{key.idx}]")
+            else:
+                parts.append(str(key))
+        return "/".join(parts)
+
+
+def get_device_mesh(
+    shape: tuple[int, ...],
+    axis_names: tuple[str, ...],
+) -> jax.sharding.Mesh:
+    """Create a Mesh object for JAX distributed arrays.
+
+    Args:
+        shape: The mesh shape, e.g., (2, 4) for 2 data-parallel x 4 model-parallel.
+        axis_names: The axis names, e.g., ("data", "model").
+
+    Returns:
+        A jax.sharding.Mesh with the specified shape and axis names.
+
+    Raises:
+        ValueError: If the product of shape does not equal the number of JAX devices.
+    """
+    import numpy as np
+
+    num_devices = len(jax.devices())
+    shape_product = math.prod(shape) if shape else 1
+
+    if shape_product != num_devices:
+        raise ValueError(
+            f"Shape product ({shape_product}) must equal number of devices "
+            f"({num_devices}). Got shape={shape}."
+        )
+
+    # Reshape devices array to match the requested shape
+    devices_array = np.array(jax.devices()).reshape(shape)
+    return jax.sharding.Mesh(devices_array, axis_names)
+
+
+def get_hardware_mesh_profile() -> dict[str, Any]:
+    """Get recommended mesh configuration for the current hardware.
+
+    Returns a dict with the following keys:
+        - device_type (str): "cpu", "gpu", or "tpu"
+        - num_devices (int): Number of available devices
+        - recommended_shape (tuple[int, ...]): Suggested mesh shape
+        - recommended_axis_names (tuple[str, ...]): Suggested axis names
+
+    This function never raises, even on unknown hardware. For a single CPU device,
+    it returns shape=(1,) and axis_names=("batch",).
+
+    Returns:
+        A dictionary with device info and recommendations.
+    """
+    try:
+        devices = jax.devices()
+        num_devices = len(devices)
+
+        # Determine device type from first device
+        device_type = devices[0].platform if hasattr(devices[0], "platform") else "cpu"
+
+        # Normalize device_type
+        if device_type.lower() in ("gpu", "cuda"):
+            device_type = "gpu"
+        elif device_type.lower() in ("tpu",):
+            device_type = "tpu"
+        else:
+            device_type = "cpu"
+
+        # Recommend shape and axis names based on device count
+        if num_devices == 1:
+            recommended_shape = (1,)
+            recommended_axis_names = ("batch",)
+        elif num_devices == 2:
+            recommended_shape = (2,)
+            recommended_axis_names = ("data",)
+        elif num_devices % 8 == 0:
+            # For TPU pods or large GPU clusters, suggest data-parallel axis
+            recommended_shape = (num_devices,)
+            recommended_axis_names = ("data",)
+        else:
+            # Default: single axis, data-parallel
+            recommended_shape = (num_devices,)
+            recommended_axis_names = ("data",)
+
+        return {
+            "device_type": device_type,
+            "num_devices": num_devices,
+            "recommended_shape": recommended_shape,
+            "recommended_axis_names": recommended_axis_names,
+        }
+    except Exception:
+        # Fallback for any errors
+        return {
+            "device_type": "cpu",
+            "num_devices": 1,
+            "recommended_shape": (1,),
+            "recommended_axis_names": ("batch",),
+        }

xtrax/engine/__init__.py

@@ -0,0 +1,6 @@
+"""Engine module for training and inference orchestration."""
+
+from xtrax.engine.engine import Engine
+from xtrax.engine.io import BoundedCallbackHandler
+
+__all__ = ["Engine", "BoundedCallbackHandler"]