xax 0.2.2__py3-none-any.whl → 0.2.4__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.2"
+__version__ = "0.2.4"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -104,6 +104,7 @@ __all__ = [
     "stage_environment",
     "to_markdown_table",
     "jit",
+    "scan",
     "save_jaxpr_dot",
     "ColoredFormatter",
     "configure_logging",
@@ -267,6 +268,7 @@ NAME_MAP: dict[str, str] = {
     "stage_environment": "utils.experiments",
     "to_markdown_table": "utils.experiments",
     "jit": "utils.jax",
+    "scan": "utils.jax",
     "save_jaxpr_dot": "utils.jaxpr",
     "ColoredFormatter": "utils.logging",
     "configure_logging": "utils.logging",
@@ -422,7 +424,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         stage_environment,
         to_markdown_table,
     )
-    from xax.utils.jax import jit
+    from xax.utils.jax import jit, scan
     from xax.utils.jaxpr import save_jaxpr_dot
     from xax.utils.logging import (
         LOG_ERROR_SUMMARY,

xax/core/state.py

@@ -1,10 +1,11 @@
 """Defines a dataclass for keeping track of the current training state."""
 
-import time
-from dataclasses import asdict, dataclass
-from typing import Any, Literal, NotRequired, TypedDict, Unpack, cast
+from dataclasses import dataclass
+from typing import Literal, NotRequired, TypedDict, Unpack, cast
 
 import jax
+import jax.numpy as jnp
+from jaxtyping import Array
 from omegaconf import MISSING
 
 from xax.core.conf import field
@@ -17,85 +18,132 @@ def _phase_to_int(phase: Phase) -> int:
 
 
 def _int_to_phase(i: int) -> Phase:
+    if i < 0 or i > 1:
+        raise ValueError(f"Invalid phase: {i}")
     return cast(Phase, ["train", "valid"][i])
 
 
 class StateDict(TypedDict, total=False):
-    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]
+    num_steps: NotRequired[int | Array]
+    num_samples: NotRequired[int | Array]
+    num_valid_steps: NotRequired[int | Array]
+    num_valid_samples: NotRequired[int | Array]
+    elapsed_time_s: NotRequired[float | Array]
+    valid_elapsed_time_s: NotRequired[float | Array]
     phase: NotRequired[Phase]
+    _phase: NotRequired[int | Array]
 
 
 @jax.tree_util.register_dataclass
 @dataclass(frozen=True, kw_only=True)
 class State:
-    num_steps: int = field(MISSING, help="Number of steps so far")
-    num_samples: int = field(MISSING, help="Number of sample so far")
-    num_valid_steps: int = field(MISSING, help="Number of validation steps so far")
-    num_valid_samples: int = field(MISSING, help="Number of validation samples so far")
-    start_time_s: float = field(MISSING, help="Start time of training")
-    elapsed_time_s: float = field(MISSING, help="Total elapsed time so far")
-    _phase: int = field(MISSING, help="Current training phase")
+    _int32_arr: Array = field(MISSING, help="Internal array for storing int64 values")
+    _float32_arr: Array = field(MISSING, help="Internal array for storing floating-point values")
+
+    @property
+    def num_steps(self) -> Array:
+        return self._int32_arr[0]
+
+    @property
+    def num_valid_steps(self) -> Array:
+        return self._int32_arr[1]
+
+    @property
+    def num_samples(self) -> Array:
+        return self._float32_arr[0]
+
+    @property
+    def num_valid_samples(self) -> Array:
+        return self._float32_arr[1]
+
+    @property
+    def elapsed_time_s(self) -> Array:
+        return self._float32_arr[2]
+
+    @property
+    def valid_elapsed_time_s(self) -> Array:
+        return self._float32_arr[3]
 
     @property
     def phase(self) -> Phase:
-        return _int_to_phase(self._phase)
+        return _int_to_phase(self._int32_arr[2].item())
 
     @classmethod
     def init_state(cls) -> "State":
         return cls(
-            num_steps=0,
-            num_samples=0,
-            num_valid_steps=0,
-            num_valid_samples=0,
-            start_time_s=time.time(),
-            elapsed_time_s=0.0,
-            _phase=0,
+            _int32_arr=jnp.array([0, 0, 0], dtype=jnp.int32),
+            _float32_arr=jnp.array([0.0, 0.0, 0.0, 0.0], dtype=jnp.float32),
         )
 
     @property
     def training(self) -> bool:
         return self.phase == "train"
 
-    def num_phase_steps(self, phase: Phase) -> int:
-        match phase:
-            case "train":
-                return self.num_steps
-            case "valid":
-                return self.num_valid_steps
-            case _:
-                raise ValueError(f"Invalid phase: {phase}")
-
     def replace(self, **kwargs: Unpack[StateDict]) -> "State":
-        extra_kwargs: dict[str, Any] = {}  # noqa: ANN401
+        int32_arr = self._int32_arr
+        float32_arr = self._float32_arr
+
+        if "num_steps" in kwargs:
+            int32_arr = int32_arr.at[0].set(kwargs["num_steps"])
+        if "num_valid_steps" in kwargs:
+            int32_arr = int32_arr.at[1].set(kwargs["num_valid_steps"])
+
         if "phase" in kwargs:
-            phase = kwargs.pop("phase")
-            match phase:
-                case "train":
-                    extra_kwargs["_phase"] = 0
-                case "valid":
-                    extra_kwargs["_phase"] = 1
-                case _:
-                    raise ValueError(f"Invalid phase: {phase}")
-        return State(**{**asdict(self), **kwargs, **extra_kwargs})
+            int32_arr = int32_arr.at[3].set(_phase_to_int(kwargs["phase"]))
+        if "_phase" in kwargs:
+            int32_arr = int32_arr.at[3].set(kwargs["_phase"])
+
+        if "num_samples" in kwargs:
+            float32_arr = float32_arr.at[0].set(kwargs["num_samples"])
+        if "num_valid_samples" in kwargs:
+            float32_arr = float32_arr.at[1].set(kwargs["num_valid_samples"])
+
+        if "elapsed_time_s" in kwargs:
+            float32_arr = float32_arr.at[2].set(kwargs["elapsed_time_s"])
+        if "valid_elapsed_time_s" in kwargs:
+            float32_arr = float32_arr.at[3].set(kwargs["valid_elapsed_time_s"])
+
+        return State(
+            _int32_arr=int32_arr,
+            _float32_arr=float32_arr,
+        )
 
     def to_dict(self) -> dict[str, int | float | str]:
         return {
             "num_steps": int(self.num_steps),
-            "num_samples": int(self.num_samples),
             "num_valid_steps": int(self.num_valid_steps),
+            "num_samples": int(self.num_samples),
             "num_valid_samples": int(self.num_valid_samples),
-            "start_time_s": float(self.start_time_s),
             "elapsed_time_s": float(self.elapsed_time_s),
+            "valid_elapsed_time_s": float(self.valid_elapsed_time_s),
             "phase": str(self.phase),
         }
 
     @classmethod
-    def from_dict(cls, d: dict[str, int | float | str]) -> "State":
+    def from_dict(cls, **d: Unpack[StateDict]) -> "State":
         if "phase" in d:
             d["_phase"] = _phase_to_int(cast(Phase, d.pop("phase")))
-        return cls(**d)  # type: ignore[arg-type]
+
+        int32_arr = jnp.array(
+            [
+                d.get("num_steps", 0),
+                d.get("num_valid_steps", 0),
+                d.get("_phase", 0),
+            ],
+            dtype=jnp.int32,
+        )
+
+        float32_arr = jnp.array(
+            [
+                d.get("num_samples", 0),
+                d.get("num_valid_samples", 0),
+                d.get("elapsed_time_s", 0.0),
+                d.get("valid_elapsed_time_s", 0.0),
+            ],
+            dtype=jnp.float32,
+        )
+
+        return cls(
+            _int32_arr=int32_arr,
+            _float32_arr=float32_arr,
+        )

xax/task/logger.py

@@ -521,7 +521,7 @@ class LoggerImpl(ABC):
         Returns:
             If the logger should log the current step.
         """
-        return self.tickers[state.phase].tick(state.elapsed_time_s)
+        return self.tickers[state.phase].tick(state.elapsed_time_s.item())
 
 
 class ToastHandler(logging.Handler):

xax/task/loggers/stdout.py

@@ -90,9 +90,9 @@ class StdoutLogger(LoggerImpl):
 
     def write_state_window(self, line: LogLine) -> None:
         state_info: dict[str, str] = {
-            "Steps": format_number(line.state.num_steps, 0),
-            "Samples": format_number(line.state.num_samples, 0),
-            "Elapsed Time": format_timedelta(datetime.timedelta(seconds=line.state.elapsed_time_s), short=True),
+            "Steps": format_number(int(line.state.num_steps.item()), 0),
+            "Samples": format_number(int(line.state.num_samples.item()), 0),
+            "Elapsed Time": format_timedelta(datetime.timedelta(seconds=line.state.elapsed_time_s.item()), short=True),
         }
 
         colored_prefix = colored("Phase: ", "grey", bold=True)

xax/task/loggers/tensorboard.py

@@ -155,14 +155,16 @@ class TensorboardLogger(LoggerImpl):
             return
 
         writer = self.get_writer(line.state.phase)
-        walltime = line.state.start_time_s + line.state.elapsed_time_s
+
+        global_step = line.state.num_steps.item()
+        walltime = line.state.elapsed_time_s.item()
 
         for namespace, scalars in line.scalars.items():
             for scalar_key, scalar_value in scalars.items():
                 writer.add_scalar(
                     f"{namespace}/{scalar_key}",
                     as_float(scalar_value.value),
-                    global_step=line.state.num_steps,
+                    global_step=global_step,
                     walltime=walltime,
                 )
 
@@ -172,7 +174,7 @@ class TensorboardLogger(LoggerImpl):
                     f"{namespace}/{distribution_key}",
                     mean=float(distribution_value.mean),
                     std=float(distribution_value.std),
-                    global_step=line.state.num_steps,
+                    global_step=global_step,
                     walltime=walltime,
                 )
 
@@ -187,7 +189,7 @@ class TensorboardLogger(LoggerImpl):
                     sum_squares=float(histogram_value.sum_squares),
                     bucket_limits=[float(x) for x in histogram_value.bucket_limits],
                     bucket_counts=[int(x) for x in histogram_value.bucket_counts],
-                    global_step=line.state.num_steps,
+                    global_step=global_step,
                     walltime=walltime,
                 )
 
@@ -196,7 +198,7 @@ class TensorboardLogger(LoggerImpl):
                 writer.add_text(
                     f"{namespace}/{string_key}",
                     string_value.value,
-                    global_step=line.state.num_steps,
+                    global_step=global_step,
                     walltime=walltime,
                 )
 
@@ -205,7 +207,7 @@ class TensorboardLogger(LoggerImpl):
                 writer.add_image(
                     f"{namespace}/{image_key}",
                     image_value.image,
-                    global_step=line.state.num_steps,
+                    global_step=global_step,
                     walltime=walltime,
                 )
 
@@ -215,7 +217,7 @@ class TensorboardLogger(LoggerImpl):
                     f"{namespace}/{video_key}",
                     video_value.frames,
                     fps=video_value.fps,
-                    global_step=line.state.num_steps,
+                    global_step=global_step,
                     walltime=walltime,
                 )
 
@@ -227,7 +229,7 @@ class TensorboardLogger(LoggerImpl):
                     faces=mesh_value.faces,
                     colors=mesh_value.colors,
                     config_dict=mesh_value.config_dict,
-                    global_step=line.state.num_steps,
+                    global_step=global_step,
                     walltime=walltime,
                 )
 

xax/task/mixins/checkpointing.py

@@ -4,7 +4,7 @@ import io
 import json
 import logging
 import tarfile
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Generic, Literal, TypeVar, cast, overload
 
@@ -76,7 +76,7 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
             if state.num_steps % self.config.save_every_n_steps == 0:
                 return True
         if self.config.save_every_n_seconds is not None:
-            last_time, cur_time = self.__last_ckpt_time, state.elapsed_time_s
+            last_time, cur_time = self.__last_ckpt_time, state.elapsed_time_s.item()
             if cur_time - last_time >= self.config.save_every_n_seconds:
                 self.__last_ckpt_time = cur_time
                 return True
@@ -200,7 +200,7 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
             def get_state() -> State:
                 if (state := tar.extractfile("state")) is None:
                     raise ValueError(f"Checkpoint does not contain a state file: {path}")
-                return State(**json.loads(state.read().decode()))
+                return State.from_dict(**json.loads(state.read().decode()))
 
             def get_config() -> Config:
                 if (config := tar.extractfile("config")) is None:
@@ -300,7 +300,7 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
                 tar.addfile(info, io.BytesIO(data))
 
             if state is not None:
-                add_file_bytes("state", json.dumps(asdict(state), indent=2).encode())
+                add_file_bytes("state", json.dumps(state.to_dict(), indent=2).encode())
             add_file_bytes("config", OmegaConf.to_yaml(self.config).encode())
 
         # Updates the symlink to the new checkpoint

xax/task/mixins/train.py

@@ -46,6 +46,7 @@ from xax.task.mixins.logger import LoggerConfig, LoggerMixin
 from xax.task.mixins.runnable import RunnableConfig, RunnableMixin
 from xax.task.mixins.step_wrapper import StepContextConfig, StepContextMixin
 from xax.utils.experiments import (
+    ContextTimer,
     StateTimer,
     TrainingFinishedError,
     diff_configs,
@@ -119,27 +120,30 @@ class ValidStepTimer:
             return True
 
         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
+            self.last_valid_time = state.elapsed_time_s.item()
+            self.last_valid_step = state.num_steps.item()
             return False
 
         # Step-based validation.
         valid_every_n_steps = self.valid_every_n_steps
         if valid_every_n_steps is not None and state.num_steps >= valid_every_n_steps + self.last_valid_step:
-            self.last_valid_step = state.num_steps
+            self.last_valid_step = state.num_steps.item()
             return True
 
         # Time-based validation.
         valid_every_n_seconds = self.valid_every_n_seconds
-        if valid_every_n_seconds is not None and state.elapsed_time_s - self.last_valid_time >= valid_every_n_seconds:
-            self.last_valid_time = state.elapsed_time_s
+        if (
+            valid_every_n_seconds is not None
+            and state.elapsed_time_s.item() - self.last_valid_time >= valid_every_n_seconds
+        ):
+            self.last_valid_time = state.elapsed_time_s.item()
             return True
 
         # Time-based validation for first validation step.
         if self.first_valid_step_flag:
             valid_first_n_seconds = self.valid_first_n_seconds
-            if valid_first_n_seconds is not None and state.elapsed_time_s >= valid_first_n_seconds:
-                self.last_valid_time = state.elapsed_time_s
+            if valid_first_n_seconds is not None and state.elapsed_time_s.item() >= valid_first_n_seconds:
+                self.last_valid_time = state.elapsed_time_s.item()
                 self.first_valid_step_flag = False
                 return True
 
@@ -214,10 +218,6 @@ class TrainMixin(
     def prng_key(self) -> PRNGKeyArray:
         return jax.random.PRNGKey(self.config.random_seed)
 
-    def on_step_end(self, state: State) -> State:
-        state = super().on_step_end(state)
-        return state.replace(elapsed_time_s=time.time() - state.start_time_s)
-
     def log_train_step(
         self,
         model: PyTree,
@@ -548,7 +548,7 @@ class TrainMixin(
             "loss": loss,
         }
 
-    @xax_jit(static_argnames=["self", "model_static"])
+    @xax_jit(static_argnames=["self", "model_static"], jit_level=3)
     def get_output_and_loss(
         self,
         model_arr: PyTree,
@@ -572,12 +572,12 @@ class TrainMixin(
         state: State,
     ) -> tuple[PyTree, optax.OptState, Output, dict[str, Array]]:
         grad_fn = jax.grad(self.get_output_and_loss, argnums=0, has_aux=True)
-        grad_fn = xax_jit(static_argnums=[1])(grad_fn)
+        grad_fn = xax_jit(static_argnums=[1], jit_level=3)(grad_fn)
         grads, (output, metrics) = grad_fn(model_arr, model_static, batch, state)
         model_arr, opt_state, grad_metrics = self.apply_gradients_with_clipping(model_arr, grads, optimizer, opt_state)
         return model_arr, opt_state, output, metrics | grad_metrics
 
-    @xax_jit(static_argnames=["self", "optimizer"])
+    @xax_jit(static_argnames=["self", "optimizer"], jit_level=3)
     def apply_gradients_with_clipping(
         self,
         model_arr: PyTree,
@@ -641,8 +641,8 @@ class TrainMixin(
     def maybe_log_termination_time(self, remaining_percent: float, state: State) -> None:
         if self._last_printed_remaining_time + PRINT_FINISH_TIME_EVERY_N_SECONDS > state.elapsed_time_s:
             return
-        self._last_printed_remaining_time = state.elapsed_time_s
-        remaining_seconds = remaining_percent * state.elapsed_time_s / (1 - remaining_percent)
+        self._last_printed_remaining_time = state.elapsed_time_s.item()
+        remaining_seconds = remaining_percent * state.elapsed_time_s.item() / (1 - remaining_percent)
         termination_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time() + remaining_seconds))
         logger.log(LOG_PING, "Estimated finish time: %s", termination_time)
 
@@ -663,11 +663,11 @@ class TrainMixin(
     def get_step(self, state: State) -> int:
         match self._step_kind:
             case "step":
-                return state.num_steps
+                return int(state.num_steps.item())
             case "sample":
-                return state.num_samples
+                return int(state.num_samples.item())
             case "second":
-                return int(state.elapsed_time_s)
+                return int(state.elapsed_time_s.item())
             case _:
                 raise ValueError(f"Invalid step kind {self._step_kind}")
 
@@ -683,7 +683,7 @@ class TrainMixin(
     def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401
         return eqx.is_inexact_array(item)
 
-    @xax_jit(static_argnames=["self", "model_static", "optimizer"])
+    @xax_jit(static_argnames=["self", "model_static", "optimizer"], jit_level=3)
     def train_step(
         self,
         model_arr: PyTree,
@@ -696,7 +696,7 @@ class TrainMixin(
         model_arr, opt_state, output, metrics = self.update(model_arr, model_static, optimizer, opt_state, batch, state)
         return model_arr, opt_state, output, FrozenDict(metrics)
 
-    @xax_jit(static_argnames=["self", "model_static"])
+    @xax_jit(static_argnames=["self", "model_static"], jit_level=3)
     def val_step(
         self,
         model_arr: PyTree,
@@ -720,34 +720,38 @@ class TrainMixin(
 
         while not self.is_training_over(state):
             if self.valid_step_timer.is_valid_step(state):
-                valid_batch = next(valid_pf)
+                with ContextTimer() as timer:
+                    valid_batch = next(valid_pf)
+                    output, metrics = self.val_step(model_arr, model_static, valid_batch, state)
+                    self.log_step(eqx.combine(model_arr, model_static), valid_batch, output, metrics, state)
+
                 state = state.replace(
                     phase="valid",
                     num_valid_steps=state.num_valid_steps + 1,
                     num_valid_samples=state.num_valid_samples + (self.get_size_of_batch(valid_batch) or 0),
+                    valid_elapsed_time_s=state.valid_elapsed_time_s + timer.elapsed_time,
                 )
 
-                output, metrics = self.val_step(model_arr, model_static, valid_batch, state)
-                self.log_step(eqx.combine(model_arr, model_static), valid_batch, output, metrics, state)
+            with ContextTimer() as timer:
+                state = self.on_step_start(state)
+                train_batch = next(train_pf)
+                model_arr, opt_state, output, metrics = self.train_step(
+                    model_arr=model_arr,
+                    model_static=model_static,
+                    optimizer=optimizer,
+                    opt_state=opt_state,
+                    batch=train_batch,
+                    state=state,
+                )
+                self.log_step(eqx.combine(model_arr, model_static), train_batch, output, metrics, state)
 
-            state = self.on_step_start(state)
-            train_batch = next(train_pf)
             state = state.replace(
                 phase="train",
                 num_steps=state.num_steps + 1,
                 num_samples=state.num_samples + (self.get_size_of_batch(train_batch) or 0),
+                elapsed_time_s=state.elapsed_time_s + timer.elapsed_time,
             )
 
-            model_arr, opt_state, output, metrics = self.train_step(
-                model_arr=model_arr,
-                model_static=model_static,
-                optimizer=optimizer,
-                opt_state=opt_state,
-                batch=train_batch,
-                state=state,
-            )
-            self.log_step(eqx.combine(model_arr, model_static), train_batch, output, metrics, state)
-
             state = self.on_step_end(state)
 
             if self.should_checkpoint(state):
@@ -843,10 +847,8 @@ class TrainMixin(
 
                 except TrainingFinishedError:
                     if is_master():
-                        show_info(
-                            f"Finished training after {state.num_steps} steps, {state.num_samples} samples",
-                            important=True,
-                        )
+                        num_steps, num_samples = int(state.num_steps), int(state.num_samples)
+                        show_info(f"Finished training after {num_steps} steps, {num_samples} samples", important=True)
                     self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
 
                 except (KeyboardInterrupt, bdb.BdbQuit):

xax/utils/experiments.py

@@ -111,8 +111,10 @@ class StateTimer:
 
     def step(self, state: State) -> None:
         cur_time = time.time()
-        self.step_timer.step(state.num_steps if state.phase == "train" else state.num_valid_steps, cur_time)
-        self.sample_timer.step(state.num_samples if state.phase == "train" else state.num_valid_samples, cur_time)
+        num_steps = int((state.num_steps if state.phase == "train" else state.num_valid_steps).item())
+        num_samples = int((state.num_samples if state.phase == "train" else state.num_valid_samples).item())
+        self.step_timer.step(num_steps, cur_time)
+        self.sample_timer.step(num_samples, cur_time)
         self.iter_timer.step(cur_time)
 
     def log_dict(self) -> dict[str, int | float | tuple[int | float, bool]]:

xax/utils/jax.py

@@ -1,5 +1,6 @@
 """Defines some utility functions for interfacing with Jax."""
 
+import functools
 import inspect
 import logging
 import os
@@ -23,6 +24,28 @@ Number = int | float | np.ndarray | jnp.ndarray
 P = ParamSpec("P")  # For function parameters
 R = TypeVar("R")  # For function return type
 
+# For control flow functions.
+Carry = TypeVar("Carry")
+X = TypeVar("X")
+Y = TypeVar("Y")
+
+
+@functools.lru_cache(maxsize=None)
+def disable_jit_level() -> int:
+    """Gets a debugging flag for disabling jitting.
+
+    For Xax's JIT'ed functions, we can set a JIT level which can be used to
+    disable jitting when we want to debug some NaN issues.
+
+    Returns:
+        The JIT level to disable.
+    """
+    return int(os.environ.get("DISABLE_JIT_LEVEL", "0"))
+
+
+def should_disable_jit(jit_level: int | None) -> bool:
+    return jit_level is not None and jit_level < disable_jit_level()
+
 
 def as_float(value: int | float | np.ndarray | jnp.ndarray) -> float:
     if isinstance(value, (int, float)):
@@ -55,6 +78,7 @@ def jit(
     inline: bool = False,
     abstracted_axes: Any | None = None,  # noqa: ANN401
     compiler_options: dict[str, Any] | None = None,
+    jit_level: int | None = None,
 ) -> Callable[[Callable[P, R]], Callable[P, R]]:
     """Wrapper function that provides utility improvements over Jax's JIT.
 
@@ -64,6 +88,8 @@ def jit(
     This is meant to be used as a decorator factory, and the decorated function
     calls `wrapped`.
     """
+    if should_disable_jit(jit_level):
+        return lambda fn: fn  # Identity function.
 
     def decorator(fn: Callable[P, R]) -> Callable[P, R]:
         class JitState:
@@ -138,3 +164,46 @@ def jit(
         return wrapped
 
     return decorator
+
+
+def scan(
+    f: Callable[[Carry, X], tuple[Carry, Y]],
+    init: Carry,
+    xs: X | None = None,
+    length: int | None = None,
+    reverse: bool = False,
+    unroll: int | bool = 1,
+    jit_level: int | None = None,
+) -> tuple[Carry, Y]:
+    """A wrapper around jax.lax.scan that allows for more flexible tracing.
+
+    If the provided JIT level is below the environment JIT level, we manually
+    unroll the scan function as a for loop.
+
+    Args:
+        f: The function to scan.
+        init: The initial value for the scan.
+        xs: The input to the scan.
+        length: The length of the scan.
+        reverse: Whether to reverse the scan.
+        unroll: The unroll factor for the scan.
+        jit_level: The JIT level to use for the scan.
+
+    Returns:
+        A tuple containing the final carry and the output of the scan.
+    """
+    if not should_disable_jit(jit_level):
+        return jax.lax.scan(f, init, xs, length, reverse, unroll)
+
+    if xs is None:
+        if length is None:
+            raise ValueError("length must be provided if xs is None")
+        xs = cast(X, [None] * length)
+
+    carry = init
+    ys = []
+    for x in cast(Iterable, xs):
+        carry, y = f(carry, x)
+        ys.append(y)
+
+    return carry, jax.tree.map(lambda *ys: jnp.stack(ys), *ys)

{xax-0.2.2.dist-info → xax-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.2.2
+Version: 0.2.4
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte

{xax-0.2.2.dist-info → xax-0.2.4.dist-info}/RECORD

@@ -1,10 +1,10 @@
-xax/__init__.py,sha256=Yj2SgoKyIAQzg3bt-hAS4gf0fqlfVBR4pv4JgpTl7-s,14182
+xax/__init__.py,sha256=bMKUtRtmVnHshkD4Ylw7ymzIPpcasJAoXnBsIdoSEng,14225
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
-xax/core/state.py,sha256=XejW1tGINYFFcNrscK8eZQsq02J7_RXa461QpmyWuLk,3337
+xax/core/state.py,sha256=bJONQ0wXgbgo1jjSqV3JtqG5tdMlli93Nax_ftZ2D0w,4552
 xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
 xax/nn/equinox.py,sha256=5fdOKRXqAVZPsV-aEez3i1wamr_oBYnG74GP1jEthjM,4843
@@ -17,7 +17,7 @@ xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/nn/ssm.py,sha256=8dLAcQ1hBaMT-kkHvwGu_ecxJeTY32WeMYmd4T4KtxA,10745
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/base.py,sha256=OnXi2hiKPGwt6ng1dutnoQSiw7lEiWFlC_vx99_JsbQ,7694
-xax/task/logger.py,sha256=peGtfnvnBKr9l6tx1V6XAsvPs0HP6ubV_aE7IJtOMNk,40868
+xax/task/logger.py,sha256=y4PGfMqKbfvPk8WCzr9MOsgG2X9E61KgeBVOYp-9kOY,40875
 xax/task/script.py,sha256=bMMIJoUtpSBvPp6-7bejTrajTXvSg0794sYLKdPIToE,972
 xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
 xax/task/launchers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -28,11 +28,11 @@ xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 xax/task/loggers/callback.py,sha256=lyuZX6Bir7xJM07ifdQIl1jlclgkiS82UO9V4y7wgPs,1582
 xax/task/loggers/json.py,sha256=_tKum6jk_gqVzO-4MqSNXbE-Mmn-yJzkRAT-N1y2zes,4139
 xax/task/loggers/state.py,sha256=6bG-NRsSUzAukYiglCT0oDj8zRMpffH4e1TKWGw1x4k,959
-xax/task/loggers/stdout.py,sha256=oeIgPkj4RyJgBuWaJK9ncLa65iBNJCWXhSF8fx3_54c,6564
-xax/task/loggers/tensorboard.py,sha256=KOL9l60tLctX-VAdNwe49H48SAJeGxph3sflJpojA-4,8337
+xax/task/loggers/stdout.py,sha256=ERLFrYe61hSSztzyxBRseobHQR72YFYjEd2i_hOeJ20,6595
+xax/task/loggers/tensorboard.py,sha256=KFlsK0zD2ubDqAXYL4Ds7NQ9F-Ke-PHwfhLOYsGcbw4,8306
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
-xax/task/mixins/checkpointing.py,sha256=2nJgqFcV-D8W-4j8TR3PvVh1g5hQUOo-_quKO-XlE4U,11398
+xax/task/mixins/checkpointing.py,sha256=8Hi-2G0EA5OFRjgiOutlk7HgkD5b-0GHazOAYxnGytM,11409
 xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 xax/task/mixins/cpu_stats.py,sha256=rO_9a82ZdsNec61ya4FpYE-rWqPhpijRSXsOfc6caFA,9595
 xax/task/mixins/data_loader.py,sha256=Tp7zqPdfH2_JuE6J6EP-fEtCQpq9MjKlGHYK7Zh-goU,6599
@@ -41,11 +41,11 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=hqDEsMp_SL6ee97iq26-G0g49OcWZZaX82JD4F22eJU,1781
 xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=v9oi9tNsNBYo-Ne_98nCG9qHX6sxvymHjsRDnL6GL-U,30871
+xax/task/mixins/train.py,sha256=XcetJ0MppV_RDhgg1M9_d9heEXo-zeN_FS3MyczeBBU,31219
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
-xax/utils/experiments.py,sha256=Hzl46_9IH5_9cKzxit-FyVUWBH-_lBs00ZciuIdnWO8,29811
-xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
+xax/utils/experiments.py,sha256=d2H63ECtVOKySMUMrQRqq4kcuZpoXqo-L931usDVAhE,29903
+xax/utils/jax.py,sha256=KQYUHjN6t6JIWa11aRSO3edcsAgTscw_dExxI6kCd9g,6767
 xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
 xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
@@ -58,8 +58,8 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.2.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.2.dist-info/METADATA,sha256=Ku0h6R6WToJ4rMYhcswGLXtIGVtzouWIGelHZFW30IM,1882
-xax-0.2.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.2.2.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.2.dist-info/RECORD,,
+xax-0.2.4.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.4.dist-info/METADATA,sha256=9hMsPCoszpjVN0rLDMlT20aYqmQwnHvl9T1V_0akl0U,1882
+xax-0.2.4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.2.4.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.4.dist-info/RECORD,,