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

xax/task/mixins/step_wrapper.py

@@ -1,53 +1,39 @@
 """Defines a mixin to wrap some steps in a context manager."""
 
+import time
 from dataclasses import dataclass
 from types import TracebackType
-from typing import ContextManager, Literal, TypeVar
+from typing import Callable, ContextManager, TypeVar
 
-import equinox as eqx
 import jax
 
 from xax.task.base import BaseConfig, BaseTask
 
-StepType = Literal[
-    "backward",
-    "change_mode",
-    "clip_grads",
-    "create_optimizers",
-    "forward",
-    "get_dataloader",
-    "get_dataset",
-    "get_prefetcher",
-    "get_model",
-    "get_optimizer",
-    "get_initial_opt_state",
-    "get_update_fn",
-    "load_checkpoint",
-    "log_losses",
-    "model_to_device",
-    "on_step_end",
-    "on_step_start",
-    "save_checkpoint",
-    "step",
-    "update_state",
-    "write_logs",
-    "zero_grads",
-]
-
 
 class StepContext(ContextManager):
     """Context manager to get the current step type."""
 
-    CURRENT_STEP: StepType | None = None
+    CURRENT_STEP: str | None = None
 
-    def __init__(self, step: StepType) -> None:
+    def __init__(
+        self,
+        step: str,
+        on_context_start: Callable[[str], None],
+        on_context_end: Callable[[str, float], None],
+    ) -> None:
         self.step = step
+        self.start_time = 0.0
+        self.on_context_start = on_context_start
+        self.on_context_end = on_context_end
 
     def __enter__(self) -> None:
         StepContext.CURRENT_STEP = self.step
+        self.start_time = time.time()
+        self.on_context_start(self.step)
 
     def __exit__(self, _t: type[BaseException] | None, _e: BaseException | None, _tr: TracebackType | None) -> None:
         StepContext.CURRENT_STEP = None
+        self.on_context_end(self.step, time.time() - self.start_time)
 
 
 @jax.tree_util.register_dataclass
@@ -63,6 +49,11 @@ 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)
+    def step_context(self, step: str) -> ContextManager:
+        return StepContext(step, self.on_context_start, self.on_context_stop)
+
+    def on_context_start(self, step: str) -> None:
+        pass
+
+    def on_context_stop(self, step: str, elapsed_time: float) -> None:
+        pass

xax/task/mixins/train.py

@@ -24,6 +24,7 @@ from typing import (
     TypeVar,
     cast,
     get_args,
+    overload,
 )
 
 import equinox as eqx
@@ -35,6 +36,7 @@ from omegaconf import DictConfig
 
 from xax.core.conf import field
 from xax.core.state import Phase, State
+from xax.nn.functions import set_random_seed
 from xax.nn.parallel import is_master
 from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
 from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin
@@ -115,7 +117,7 @@ class ValidStepTimer:
         if self.last_valid_time is None or self.last_valid_step is None:
             self.last_valid_time = state.elapsed_time_s
             self.last_valid_step = state.num_steps
-            return True
+            return False
 
         # Step-based validation.
         valid_every_n_steps = self.valid_every_n_steps
@@ -154,7 +156,6 @@ class TrainConfig(
     valid_first_n_steps: int = field(0, help="Treat the first N steps as validation steps")
     valid_every_n_seconds: float | None = field(60.0 * 10.0, help="Run validation every N seconds")
     valid_first_n_seconds: float | None = field(60.0, help="Run first validation after N seconds")
-    batch_dim: int = field(0, help="The batch dimension, for splitting batches into chunks")
     max_steps: int | None = field(None, help="Maximum number of steps to run")
     step_kind: str = field("step", help=f"How to measure a step; one of [{', '.join(get_args(StepKind))}]")
     random_seed: int = field(1337, help="Random seed for the task")
@@ -183,6 +184,9 @@ class TrainMixin(
     def __init__(self, config: Config) -> None:
         super().__init__(config)
 
+        # Sets the random seed whenever we instantiate a new train mixin.
+        set_random_seed(self.config.random_seed)
+
         # Timer for validation steps.
         self.valid_step_timer = ValidStepTimer(
             valid_every_n_steps=config.valid_every_n_steps,
@@ -279,31 +283,53 @@ class TrainMixin(
     def get_initial_opt_state(self, model: PyTree, optimizer: optax.GradientTransformation) -> optax.OptState:
         return optimizer.init(eqx.filter(model, eqx.is_array))
 
+    @overload
+    def load_initial_state(
+        self,
+        key: PRNGKeyArray,
+        load_optimizer: Literal[False] = False,
+    ) -> tuple[PyTree, State]: ...
+
+    @overload
     def load_initial_state(
         self,
         key: PRNGKeyArray,
-    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State]:
+        load_optimizer: Literal[True],
+    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State]: ...
+
+    def load_initial_state(
+        self,
+        key: PRNGKeyArray,
+        load_optimizer: bool = False,
+    ) -> tuple[PyTree, State] | tuple[PyTree, optax.GradientTransformation, optax.OptState, State]:
         init_ckpt_path = self.get_init_ckpt_path()
 
         if init_ckpt_path is not None:
             logger.info("Loading checkpoint from %s", init_ckpt_path)
-            with self.step_context("load_checkpoint"):
+            if load_optimizer:
                 model, optimizer, opt_state, state, config = self.load_checkpoint(init_ckpt_path)
                 config_diff = get_diff_string(diff_configs(config, cast(DictConfig, self.config)))
                 if config_diff:
                     logger.warning("Loaded config differs from current config:\n%s", config_diff)
                 return model, optimizer, opt_state, state
 
-        with self.step_context("get_model"):
-            model = self.get_model(key)
+            else:
+                model, state, config = self.load_checkpoint(init_ckpt_path, "model_state_config")
+                config_diff = get_diff_string(diff_configs(config, cast(DictConfig, self.config)))
+                if config_diff:
+                    logger.warning("Loaded config differs from current config:\n%s", config_diff)
+                return model, state
+
+        model = self.get_model(key)
+        state = State.init_state()
 
-        with self.step_context("get_optimizer"):
-            optimizer = self.get_optimizer()
+        if not load_optimizer:
+            return model, state
 
-        with self.step_context("get_initial_opt_state"):
-            opt_state = self.get_initial_opt_state(model, optimizer)
+        optimizer = self.get_optimizer()
+        opt_state = self.get_initial_opt_state(model, optimizer)
 
-        return model, optimizer, opt_state, State.init_state()
+        return model, optimizer, opt_state, state
 
     @eqx.filter_jit
     def get_output(self, model: PyTree, batch: Batch) -> Output:
@@ -424,6 +450,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, "JAX devices: %s", jax.devices())
         self.logger.log_file("git_state.txt", get_git_state(self))
         self.logger.log_file("training_code.txt", get_training_code(self))
         self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
@@ -456,7 +483,8 @@ class TrainMixin(
         while not self.is_training_over(state):
             if self.valid_step_timer.is_valid_step(state):
                 valid_batch = next(valid_pf)
-                model, loss, output = self.val_step(model, valid_batch)
+                with self.step_context("model_step"):
+                    model, loss, output = self.val_step(model, valid_batch)
 
                 # Perform logging.
                 with self.step_context("write_logs"):
@@ -464,22 +492,19 @@ class TrainMixin(
                     self.log_step(model, valid_batch, output, loss, state)
                     state.num_valid_samples += 1
 
-            with self.step_context("on_step_start"):
-                state = self.on_step_start(state)
+            state = self.on_step_start(state)
 
-            with self.step_context("update_state"):
+            with self.step_context("model_step"):
                 train_batch = next(train_pf)
                 model, opt_state, loss, output = self.train_step(model, optimizer, opt_state, train_batch)
 
-            # Perform logging.
             with self.step_context("write_logs"):
                 state.phase = "train"
                 self.log_step(model, train_batch, output, loss, state)
                 state.num_steps += 1
                 state.num_samples += self.get_size_of_batch(train_batch) or 0
 
-            with self.step_context("on_step_end"):
-                state = self.on_step_end(state)
+            state = self.on_step_end(state)
 
             if self.should_checkpoint(state):
                 self.save_checkpoint(model, optimizer, opt_state, state)
@@ -496,14 +521,9 @@ class TrainMixin(
         except NotImplementedError:
             pass
 
-        with self.step_context("get_dataset"):
-            train_ds = self.get_dataset("train")
-
-        with self.step_context("get_dataloader"):
-            train_dl = self.get_dataloader(train_ds, "train")
-
-        with self.step_context("get_prefetcher"):
-            train_pf = self.get_prefetcher(train_dl)
+        train_ds = self.get_dataset("train")
+        train_dl = self.get_dataloader(train_ds, "train")
+        train_pf = self.get_prefetcher(train_dl)
 
         try:
             with train_pf as train_pf_ctx:
@@ -520,14 +540,9 @@ class TrainMixin(
         except NotImplementedError:
             pass
 
-        with self.step_context("get_dataset"):
-            valid_ds = self.get_dataset("valid")
-
-        with self.step_context("get_dataloader"):
-            valid_dl = self.get_dataloader(valid_ds, "valid")
-
-        with self.step_context("get_prefetcher"):
-            valid_pf = self.get_prefetcher(valid_dl)
+        valid_ds = self.get_dataset("valid")
+        valid_dl = self.get_dataloader(valid_ds, "valid")
+        valid_pf = self.get_prefetcher(valid_dl)
 
         try:
             with valid_pf as valid_pf_ctx:
@@ -559,7 +574,7 @@ class TrainMixin(
                 Thread(target=self.log_state, daemon=True).start()
 
             key, model_key = jax.random.split(key)
-            model, optimizer, opt_state, state = self.load_initial_state(model_key)
+            model, optimizer, opt_state, state = self.load_initial_state(model_key, load_optimizer=True)
             state = self.on_training_start(state)
 
             def on_exit() -> None:

xax/task/script.py

@@ -17,8 +17,6 @@ from xax.task.mixins import (
     ProcessMixin,
     RunnableConfig,
     RunnableMixin,
-    StepContextConfig,
-    StepContextMixin,
 )
 
 
@@ -28,7 +26,6 @@ class ScriptConfig(
     GPUStatsConfig,
     ProcessConfig,
     LoggerConfig,
-    StepContextConfig,
     ArtifactsConfig,
     RunnableConfig,
     BaseConfig,
@@ -44,7 +41,6 @@ class Script(
     GPUStatsMixin[ConfigT],
     ProcessMixin[ConfigT],
     LoggerMixin[ConfigT],
-    StepContextMixin[ConfigT],
     ArtifactsMixin[ConfigT],
     RunnableMixin[ConfigT],
     BaseTask[ConfigT],

xax/utils/debugging.py

@@ -0,0 +1,49 @@
+"""Defines some useful Jax debugging utilities."""
+
+from collections import deque
+from collections.abc import Iterable, Mapping
+from typing import Any, Callable, Deque
+
+from jaxtyping import Array
+
+
+def get_named_leaves(
+    obj: Any,  # noqa: ANN401
+    is_leaf: Callable[[Any], bool] = lambda x: isinstance(x, Array),  # noqa: ANN401
+    max_depth: int = 100,
+) -> list[tuple[str, Any]]:  # noqa: ANN401
+    ret: list[tuple[str, Any]] = []
+    q: Deque[tuple[int, str, Any]] = deque()  # noqa: ANN401
+    q.append((0, "", obj))
+
+    while q:
+        depth, name, node = q.popleft()
+
+        if depth > max_depth:
+            continue
+
+        if hasattr(node, "__dict__") and isinstance(node.__dict__, Mapping):
+            for cname, cnode in node.__dict__.items():
+                gname = f"{name}.{cname}" if name else cname
+                if is_leaf(cnode):
+                    ret.append((gname, cnode))
+                else:
+                    q.append((depth + 1, gname, cnode))
+
+        elif isinstance(node, Mapping):
+            for cname, cnode in node.items():
+                gname = f"{name}.{cname}" if name else cname
+                if is_leaf(cnode):
+                    ret.append((gname, cnode))
+                else:
+                    q.append((depth + 1, gname, cnode))
+
+        elif isinstance(node, Iterable):
+            for i, cnode in enumerate(node):
+                gname = f"{name}.{i}" if name else str(i)
+                if is_leaf(cnode):
+                    ret.append((gname, cnode))
+                else:
+                    q.append((depth + 1, gname, cnode))
+
+    return ret

xax/utils/experiments.py

@@ -23,7 +23,8 @@ import urllib.request
 import warnings
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Iterator, TypeVar, cast
+from types import TracebackType
+from typing import Any, Iterator, Self, TypeVar, cast
 from urllib.parse import urlparse
 
 import git
@@ -116,19 +117,19 @@ class StateTimer:
         logs: dict[str, dict[str, int | float]] = {}
 
         # Logs step statistics.
-        logs["⏰ steps"] = {
+        logs["⌛ steps"] = {
             "total": self.step_timer.steps,
             "per-second": self.step_timer.steps_per_second,
         }
 
         # Logs sample statistics.
-        logs["⏰ samples"] = {
+        logs["⌛ samples"] = {
             "total": self.sample_timer.steps,
             "per-second": self.sample_timer.steps_per_second,
         }
 
         # Logs full iteration statistics.
-        logs["🔧 dt"] = {
+        logs["⌛ dt"] = {
             "iter": self.iter_timer.iter_seconds,
         }
 
@@ -147,6 +148,24 @@ class IntervalTicker:
         return False
 
 
+class ContextTimer:
+    def __init__(self) -> None:
+        self.start_time = 0.0
+        self.elapsed_time = 0.0
+
+    def __enter__(self) -> Self:
+        self.start_time = time.time()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.elapsed_time = time.time() - self.start_time
+
+
 def abs_path(path: str) -> str:
     return str(Path(path).resolve())
 

xax/utils/jaxpr.py

@@ -0,0 +1,77 @@
+"""Visualize JAXPR."""
+
+from pathlib import Path
+
+import jax
+import jax.core
+
+
+def save_jaxpr_dot(closed_jaxpr: jax.core.ClosedJaxpr, filename: str | Path) -> None:
+    """Save the JAXPR to a DOT file.
+
+    Example usage:
+
+        grad_fn_jaxpr = jax.make_jaxpr(loss_fn)(variables)
+        save_jaxpr_dot(grad_fn_jaxpr, "grad_fn_jaxpr.dot")
+
+    Then, you can visualize the JAXPR using Graphviz:
+
+        dot -Tpng grad_fn_jaxpr.dot > grad_fn_jaxpr.png
+
+    Args:
+        closed_jaxpr: The closed JAXPR to save.
+        filename: The filename to save the JAXPR to.
+    """
+    if hasattr(closed_jaxpr, "jaxpr"):
+        jaxpr = closed_jaxpr.jaxpr
+    else:
+        jaxpr = closed_jaxpr
+
+    with open(filename, "w") as f:
+        f.write("digraph Jaxpr {\n")
+
+        var_names: dict[jax.core.Var, str] = {}
+        var_count = 0
+
+        def get_var_name(var: jax.core.Var) -> str:
+            """Get a unique name for a variable."""
+            nonlocal var_names, var_count
+
+            # Handle Literal objects specially since they're not hashable
+            if isinstance(var, jax.core.Literal):
+                # Create a name based on the literal value
+                name = f"lit_{var.val}"
+                return name
+
+            # For other variables
+            if var not in var_names:
+                name = f"var_{var_count}"
+                var_names[var] = name
+                var_count += 1
+            return var_names[var]
+
+        for var in jaxpr.invars:
+            node_name = get_var_name(var)
+            f.write(f'  {node_name} [label="{node_name}\\n(input)"];\n')
+
+        eq_count = 0
+        for eq in jaxpr.eqns:
+            eq_node = f"eq{eq_count}"
+            label = f"{eq.primitive.name}"
+            f.write(f'  {eq_node} [shape=box, label="{label}"];\n')
+
+            for invar in eq.invars:
+                var_name = get_var_name(invar)
+                f.write(f"  {var_name} -> {eq_node};\n")
+
+            for outvar in eq.outvars:
+                var_name = get_var_name(outvar)
+                f.write(f"  {eq_node} -> {var_name};\n")
+
+            eq_count += 1
+
+        for var in jaxpr.outvars:
+            node_name = get_var_name(var)
+            f.write(f'  {node_name} [peripheries=2, label="{node_name}\\n(output)"];\n')
+
+        f.write("}\n")

xax/utils/logging.py

@@ -140,7 +140,13 @@ class ColoredFormatter(logging.Formatter):
         return logging.Formatter.format(self, record)
 
 
-def configure_logging(prefix: str | None = None, *, rank: int | None = None, world_size: int | None = None) -> None:
+def configure_logging(
+    prefix: str | None = None,
+    *,
+    rank: int | None = None,
+    world_size: int | None = None,
+    debug: bool | None = None,
+) -> None:
     """Instantiates logging.
 
     This captures logs and reroutes them to the Toasts module, which is
@@ -151,6 +157,7 @@ def configure_logging(prefix: str | None = None, *, rank: int | None = None, wor
         prefix: An optional prefix to add to the logger
         rank: The current rank, or None if not using multiprocessing
         world_size: The total world size, or None if not using multiprocessing
+        debug: Whether to enable debug logging
     """
     if rank is not None or world_size is not None:
         assert rank is not None and world_size is not None
@@ -168,7 +175,10 @@ def configure_logging(prefix: str | None = None, *, rank: int | None = None, wor
     stream_handler.addFilter(filter)
     root_logger.addHandler(stream_handler)
 
-    root_logger.setLevel(logging._nameToLevel[config.log_level])
+    if debug is None:
+        root_logger.setLevel(logging._nameToLevel[config.log_level])
+    else:
+        root_logger.setLevel(logging.DEBUG if debug else logging.INFO)
 
     # Avoid junk logs from other libraries.
     if config.hide_third_party_logs: