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

xax/task/mixins/checkpointing.py

@@ -9,6 +9,7 @@ from pathlib import Path
 from typing import Any, Callable, Generic, Literal, TypeVar, cast, overload
 
 import cloudpickle
+import jax
 import optax
 from jaxtyping import PyTree
 from omegaconf import DictConfig, OmegaConf
@@ -38,6 +39,7 @@ def get_ckpt_path(exp_dir: Path, state: State | None = None) -> Path:
     return exp_dir / "checkpoints" / f"ckpt.{state.num_steps}.bin"
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class CheckpointingConfig(ArtifactsConfig):
     save_every_n_steps: int | None = field(None, help="Save a checkpoint every N steps")
@@ -45,6 +47,7 @@ class CheckpointingConfig(ArtifactsConfig):
     only_save_most_recent: bool = field(True, help="Only keep the most recent checkpoint")
     load_from_ckpt_path: str | None = field(None, help="If set, load initial model weights from this path")
     load_ckpt_strict: bool = field(True, help="If set, only load weights for which have a matching key in the model")
+    save_tf_model: bool = field(False, help="If set, saves a Tensorflow version of the model")
 
 
 Config = TypeVar("Config", bound=CheckpointingConfig)
@@ -196,6 +199,15 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
             add_file("state", lambda buf: buf.write(json.dumps(asdict(state), indent=2).encode()))
             add_file("config", lambda buf: buf.write(OmegaConf.to_yaml(self.config).encode()))
 
+        if self.config.save_tf_model:
+            try:
+                from jax.experimental import jax2tf
+            except ModuleNotFoundError:
+                raise ImportError("Tensorflow is not installed. Install it with `pip install tensorflow`")
+
+            tf_model = jax2tf.convert(model)
+            add_file("model.tf", lambda buf: cloudpickle.dump(tf_model, buf))
+
         # Updates the symlink to the new checkpoint.
         last_ckpt_path.unlink(missing_ok=True)
         try:

xax/task/mixins/compile.py

@@ -0,0 +1,104 @@
+"""Defines a mixin for handling JAX compilation behavior.
+
+This mixin allows control over JAX compilation settings like jit, pmap, and vmap
+behavior during initialization and training.
+"""
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Generic, TypeVar
+
+import jax
+
+from xax.core.conf import field
+from xax.task.base import BaseConfig, BaseTask
+
+logger = logging.getLogger(__name__)
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class CompileOptions:
+    # JAX compilation options
+    disable_jit: bool = field(
+        value=False,
+        help="If True, disables JIT compilation",
+    )
+    enable_x64: bool = field(
+        value=False,
+        help="If True, enables 64-bit precision",
+    )
+    default_device: str | None = field(
+        value=None,
+        help="Default device to use (e.g. 'cpu', 'gpu')",
+    )
+
+    # JAX logging options
+    logging_level: str = field(
+        value="INFO",
+        help="JAX logging verbosity level",
+    )
+
+    # JAX cache options
+    cache_dir: str | None = field(
+        value=lambda: str((Path.home() / ".cache" / "jax" / "jaxcache").resolve()),
+        help="Directory for JAX compilation cache. If None, caching is disabled",
+    )
+    cache_min_size_bytes: int = field(
+        value=-1,
+        help="Minimum size in bytes for cache entries. -1 means no minimum",
+    )
+    cache_min_compile_time_secs: float = field(
+        value=0.0,
+        help="Minimum compilation time in seconds for cache entries. 0 means no minimum",
+    )
+    cache_enable_xla: str = field(
+        value="all",
+        help="Which XLA caches to enable",
+    )
+
+
+@jax.tree_util.register_dataclass
+@dataclass
+class CompileConfig(BaseConfig):
+    compile: CompileOptions = field(CompileOptions(), help="Compilation configuration")
+
+
+Config = TypeVar("Config", bound=CompileConfig)
+
+
+class CompileMixin(BaseTask[Config], Generic[Config]):
+    """Defines a task mixin for controlling JAX compilation behavior."""
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        cc = self.config.compile
+
+        # Set basic compilation flags
+        if cc.disable_jit:
+            logger.info("Disabling JIT compilation")
+            jax.config.update("jax_disable_jit", True)
+
+        if cc.enable_x64:
+            logger.info("Enabling 64-bit precision")
+            jax.config.update("jax_enable_x64", True)
+
+        if cc.default_device is not None:
+            logger.info("Setting default device to %s", cc.default_device)
+            jax.config.update("jax_default_device", cc.default_device)
+
+        # Set logging level
+        logger.info("Setting JAX logging level to %s", cc.logging_level)
+        jax.config.update("jax_logging_level", cc.logging_level)
+
+        # Configure compilation cache
+        if cc.cache_dir is not None:
+            logger.info("Setting JAX compilation cache directory to %s", cc.cache_dir)
+            jax.config.update("jax_compilation_cache_dir", cc.cache_dir)
+
+            logger.info("Configuring JAX compilation cache parameters")
+            jax.config.update("jax_persistent_cache_min_entry_size_bytes", cc.cache_min_size_bytes)
+            jax.config.update("jax_persistent_cache_min_compile_time_secs", cc.cache_min_compile_time_secs)
+            jax.config.update("jax_persistent_cache_enable_xla_caches", cc.cache_enable_xla)

xax/task/mixins/cpu_stats.py

@@ -6,15 +6,16 @@ leaks in your dataloader, among other issues.
 """
 
 import logging
-import multiprocessing as mp
 import os
 import time
 from ctypes import Structure, c_double, c_uint16, c_uint64
 from dataclasses import dataclass
+from multiprocessing.context import BaseContext, Process
 from multiprocessing.managers import SyncManager, ValueProxy
 from multiprocessing.synchronize import Event
 from typing import Generic, TypeVar
 
+import jax
 import psutil
 
 from xax.core.conf import field
@@ -26,12 +27,14 @@ from xax.task.mixins.process import ProcessConfig, ProcessMixin
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class CPUStatsOptions:
     ping_interval: int = field(1, help="How often to check stats (in seconds)")
     only_log_once: bool = field(False, help="If set, only log read stats one time")
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class CPUStatsConfig(ProcessConfig, LoggerConfig, BaseConfig):
     cpu_stats: CPUStatsOptions = field(CPUStatsOptions(), help="CPU stats configuration")
@@ -55,7 +58,7 @@ class CPUStats(Structure):
     ]
 
 
-@dataclass
+@dataclass(kw_only=True)
 class CPUStatsInfo:
     cpu_percent: float
     mem_percent: float
@@ -142,9 +145,16 @@ def worker(
 
 
 class CPUStatsMonitor:
-    def __init__(self, ping_interval: float, manager: SyncManager) -> None:
+    def __init__(
+        self,
+        ping_interval: float,
+        context: BaseContext,
+        manager: SyncManager,
+    ) -> None:
         self._ping_interval = ping_interval
         self._manager = manager
+        self._context = context
+
         self._monitor_event = self._manager.Event()
         self._start_event = self._manager.Event()
         self._cpu_stats_smem = self._manager.Value(
@@ -163,7 +173,7 @@ class CPUStatsMonitor:
             ),
         )
         self._cpu_stats: CPUStatsInfo | None = None
-        self._proc: mp.Process | None = None
+        self._proc: Process | None = None
 
     def get_if_set(self) -> CPUStatsInfo | None:
         if self._monitor_event.is_set():
@@ -184,7 +194,7 @@ class CPUStatsMonitor:
         if self._start_event.is_set():
             self._start_event.clear()
         self._cpu_stats = None
-        self._proc = mp.Process(
+        self._proc = self._context.Process(  # type: ignore[attr-defined]
             target=worker,
             args=(self._ping_interval, self._cpu_stats_smem, self._monitor_event, self._start_event, os.getpid()),
             daemon=True,
@@ -215,6 +225,7 @@ class CPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
 
         self._cpu_stats_monitor = CPUStatsMonitor(
             ping_interval=self.config.cpu_stats.ping_interval,
+            context=self._mp_ctx,
             manager=self._mp_manager,
         )
 

xax/task/mixins/data_loader.py

@@ -1,9 +1,9 @@
 """Defines a mixin for instantiating dataloaders."""
 
 import logging
-from abc import ABC, abstractmethod
+from abc import ABC
 from dataclasses import dataclass
-from typing import Generic, TypeVar
+from typing import Generic, Iterator, TypeVar
 
 import jax
 from dpshdl.dataloader import CollatedDataloaderItem, Dataloader
@@ -13,7 +13,7 @@ from omegaconf import II, MISSING
 
 from xax.core.conf import field, is_missing
 from xax.core.state import Phase
-from xax.nn.functions import recursive_apply, set_random_seed
+from xax.nn.functions import set_random_seed
 from xax.task.base import BaseConfig, BaseTask
 from xax.task.mixins.process import ProcessConfig, ProcessMixin
 from xax.utils.logging import LOG_ERROR_SUMMARY, configure_logging
@@ -24,6 +24,7 @@ T = TypeVar("T")
 Tc_co = TypeVar("Tc_co", covariant=True)
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class DataloaderErrorConfig:
     sleep_backoff: float = field(0.1, help="The initial sleep time after an exception")
@@ -36,6 +37,7 @@ class DataloaderErrorConfig:
     log_exceptions_all_workers: bool = field(False, help="If set, log exceptions from all workers")
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class DataloaderConfig:
     num_workers: int | None = field(MISSING, help="Number of workers for loading samples")
@@ -43,6 +45,7 @@ class DataloaderConfig:
     error: DataloaderErrorConfig = field(DataloaderErrorConfig(), help="Dataloader error configuration")
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class DataloadersConfig(ProcessConfig, BaseConfig):
     batch_size: int = field(MISSING, help="Size of each batch")
@@ -63,9 +66,6 @@ Config = TypeVar("Config", bound=DataloadersConfig)
 
 class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config], ABC):
     def __init__(self, config: Config) -> None:
-        if is_missing(config, "batch_size"):
-            config.batch_size = self.get_batch_size()
-
         super().__init__(config)
 
     def get_batch_size(self) -> int:
@@ -74,6 +74,12 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
             "method to return the desired training batch size."
         )
 
+    @property
+    def batch_size(self) -> int:
+        if is_missing(self.config, "batch_size"):
+            self.config.batch_size = self.get_batch_size()
+        return self.config.batch_size
+
     def dataloader_config(self, phase: Phase) -> DataloaderConfig:
         match phase:
             case "train":
@@ -83,7 +89,6 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
-    @abstractmethod
     def get_dataset(self, phase: Phase) -> Dataset:
         """Returns the dataset for the given phase.
 
@@ -93,6 +98,16 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
         Returns:
             The dataset for the given phase.
         """
+        raise NotImplementedError(
+            "You must implement either the `get_dataset` method to return the dataset for the given phase, "
+            "or `get_data_iterator` to return an iterator for the given dataset."
+        )
+
+    def get_data_iterator(self, phase: Phase) -> Iterator:
+        raise NotImplementedError(
+            "You must implement either the `get_dataset` method to return the dataset for the given phase, "
+            "or `get_data_iterator` to return an iterator for the given dataset."
+        )
 
     def get_dataloader(self, dataset: Dataset[T, Tc_co], phase: Phase) -> Dataloader[T, Tc_co]:
         debugging = self.config.debug_dataloader
@@ -128,11 +143,7 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
         )
 
     def get_prefetcher(self, dataloader: Dataloader[T, Tc_co]) -> Prefetcher[Tc_co, Tc_co]:
-        return Prefetcher(to_device_func=self.to_device_fn, dataloader=dataloader)
-
-    @classmethod
-    def to_device_fn(cls, sample: T) -> T:
-        return recursive_apply(sample, jax.device_put, include_numpy=True)
+        return Prefetcher(to_device_func=jax.device_put, dataloader=dataloader)
 
     @classmethod
     def dataloader_worker_init_fn(cls, worker_id: int, num_workers: int) -> None:

xax/task/mixins/gpu_stats.py

@@ -6,17 +6,19 @@ This logs GPU memory and utilization in a background process using
 
 import functools
 import logging
-import multiprocessing as mp
 import os
 import re
 import shutil
 import subprocess
 from ctypes import Structure, c_double, c_uint32
 from dataclasses import dataclass
+from multiprocessing.context import BaseContext, Process
 from multiprocessing.managers import SyncManager, ValueProxy
 from multiprocessing.synchronize import Event
 from typing import Generic, Iterable, Pattern, TypeVar
 
+import jax
+
 from xax.core.conf import field
 from xax.core.state import State
 from xax.task.mixins.logger import LoggerConfig, LoggerMixin
@@ -25,12 +27,14 @@ from xax.task.mixins.process import ProcessConfig, ProcessMixin
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class GPUStatsOptions:
     ping_interval: int = field(10, help="How often to check stats (in seconds)")
     only_log_once: bool = field(False, help="If set, only log read stats one time")
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class GPUStatsConfig(ProcessConfig, LoggerConfig):
     gpu_stats: GPUStatsOptions = field(GPUStatsOptions(), help="GPU stats configuration")
@@ -147,8 +151,14 @@ def worker(
 
 
 class GPUStatsMonitor:
-    def __init__(self, ping_interval: float, manager: SyncManager) -> None:
+    def __init__(
+        self,
+        ping_interval: float,
+        context: BaseContext,
+        manager: SyncManager,
+    ) -> None:
         self._ping_interval = ping_interval
+        self._context = context
         self._manager = manager
 
         num_gpus = get_num_gpus()
@@ -169,7 +179,7 @@ class GPUStatsMonitor:
             for i in range(num_gpus)
         ]
         self._gpu_stats: dict[int, GPUStatsInfo] = {}
-        self._proc: mp.Process | None = None
+        self._proc: Process | None = None
 
     def get_if_set(self) -> dict[int, GPUStatsInfo]:
         gpu_stats: dict[int, GPUStatsInfo] = {}
@@ -196,7 +206,7 @@ class GPUStatsMonitor:
         if self._start_event.is_set():
             self._start_event.clear()
         self._gpu_stats.clear()
-        self._proc = mp.Process(
+        self._proc = self._context.Process(  # type: ignore[attr-defined]
             target=worker,
             args=(self._ping_interval, self._smems, self._main_event, self._events, self._start_event),
             daemon=True,
@@ -226,7 +236,11 @@ class GPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
 
         self._gpu_stats_monitor = None
         if shutil.which("nvidia-smi") is not None:
-            self._gpu_stats_monitor = GPUStatsMonitor(config.gpu_stats.ping_interval, self._mp_manager)
+            self._gpu_stats_monitor = GPUStatsMonitor(
+                config.gpu_stats.ping_interval,
+                self._mp_ctx,
+                self._mp_manager,
+            )
 
     def on_training_start(self, state: State) -> State:
         state = super().on_training_start(state)

xax/task/mixins/logger.py

@@ -6,7 +6,8 @@ from pathlib import Path
 from types import TracebackType
 from typing import Generic, Self, TypeVar
 
-from xax.core.conf import Device as BaseDeviceConfig, field
+import jax
+
 from xax.core.state import State
 from xax.task.base import BaseConfig, BaseTask
 from xax.task.logger import Logger, LoggerImpl
@@ -18,9 +19,10 @@ from xax.task.mixins.artifacts import ArtifactsMixin
 from xax.utils.text import is_interactive_session
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class LoggerConfig(BaseConfig):
-    device: BaseDeviceConfig = field(BaseDeviceConfig(), help="Device configuration")
+    pass
 
 
 Config = TypeVar("Config", bound=LoggerConfig)

xax/task/mixins/process.py

@@ -7,6 +7,8 @@ from multiprocessing.context import BaseContext
 from multiprocessing.managers import SyncManager
 from typing import Generic, TypeVar
 
+import jax
+
 from xax.core.conf import field
 from xax.core.state import State
 from xax.task.base import BaseConfig, BaseTask
@@ -14,9 +16,10 @@ from xax.task.base import BaseConfig, BaseTask
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class ProcessConfig(BaseConfig):
-    multiprocessing_context: str | None = field(None, help="The multiprocessing context to use")
+    multiprocessing_context: str | None = field("spawn", help="The multiprocessing context to use")
 
 
 Config = TypeVar("Config", bound=ProcessConfig)

xax/task/mixins/runnable.py

@@ -6,10 +6,13 @@ from dataclasses import dataclass
 from types import FrameType
 from typing import Callable, TypeVar
 
+import jax
+
 from xax.task.base import BaseConfig, BaseTask, RawConfigType
 from xax.task.launchers.base import BaseLauncher
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class RunnableConfig(BaseConfig):
     pass

xax/task/mixins/step_wrapper.py

@@ -4,6 +4,9 @@ from dataclasses import dataclass
 from types import TracebackType
 from typing import ContextManager, Literal, TypeVar
 
+import equinox as eqx
+import jax
+
 from xax.task.base import BaseConfig, BaseTask
 
 StepType = Literal[
@@ -47,6 +50,7 @@ class StepContext(ContextManager):
         StepContext.CURRENT_STEP = None
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class StepContextConfig(BaseConfig):
     pass
@@ -59,5 +63,6 @@ class StepContextMixin(BaseTask[Config]):
     def __init__(self, config: Config) -> None:
         super().__init__(config)
 
+    @eqx.filter_jit
     def step_context(self, step: StepType) -> ContextManager:
         return StepContext(step)