xax 0.2.6__py3-none-any.whl → 0.2.7__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.6"
+__version__ = "0.2.7"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -66,11 +66,13 @@ __all__ = [
     "StateLogger",
     "StdoutLogger",
     "TensorboardLogger",
+    "load_ckpt",
     "CPUStatsOptions",
     "DataloaderConfig",
     "GPUStatsOptions",
     "StepContext",
     "ValidStepTimer",
+    "get_param_count",
     "Script",
     "ScriptConfig",
     "Config",
@@ -230,11 +232,13 @@ NAME_MAP: dict[str, str] = {
     "StateLogger": "task.loggers.state",
     "StdoutLogger": "task.loggers.stdout",
     "TensorboardLogger": "task.loggers.tensorboard",
+    "load_ckpt": "task.mixins.checkpointing",
     "CPUStatsOptions": "task.mixins.cpu_stats",
     "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
     "StepContext": "task.mixins.step_wrapper",
     "ValidStepTimer": "task.mixins.train",
+    "get_param_count": "task.mixins.train",
     "Script": "task.script",
     "ScriptConfig": "task.script",
     "Config": "task.task",
@@ -390,11 +394,12 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.loggers.state import StateLogger
     from xax.task.loggers.stdout import StdoutLogger
     from xax.task.loggers.tensorboard import TensorboardLogger
+    from xax.task.mixins.checkpointing import load_ckpt
     from xax.task.mixins.cpu_stats import CPUStatsOptions
     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, Output, ValidStepTimer, get_param_count
     from xax.task.script import Script, ScriptConfig
     from xax.task.task import Config, Task
     from xax.utils.data.collate import CollateMode, collate, collate_non_null

xax/nn/functions.py

@@ -1,5 +1,5 @@
 # mypy: disable-error-code="override"
-"""Defines helper Torch functions."""
+"""Defines helper Jax functions."""
 
 import random
 from dataclasses import is_dataclass

xax/task/loggers/json.py

@@ -2,7 +2,6 @@
 
 import json
 import sys
-from dataclasses import asdict
 from typing import Any, Literal, Mapping, TextIO
 
 from jaxtyping import Array
@@ -67,7 +66,7 @@ class JsonLogger(LoggerImpl):
         return self.err_log_stream
 
     def get_json(self, line: LogLine) -> str:
-        data: dict = {"state": asdict(line.state)}
+        data: dict = {"state": line.state.to_dict()}
 
         def add_logs(log: Mapping[str, Mapping[str, LogScalar | LogString]], data: dict) -> None:
             for namespace, values in log.items():

xax/task/mixins/checkpointing.py

@@ -52,6 +52,114 @@ class CheckpointingConfig(ArtifactsConfig):
 Config = TypeVar("Config", bound=CheckpointingConfig)
 
 
+@overload
+def load_ckpt(
+    path: Path,
+    *,
+    part: Literal["all"],
+    model_template: PyTree,
+    optimizer_template: PyTree,
+    opt_state_template: PyTree,
+) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State, DictConfig]: ...
+
+
+@overload
+def load_ckpt(
+    path: Path,
+    *,
+    part: Literal["model_state_config"],
+    model_template: PyTree,
+) -> tuple[PyTree, State, DictConfig]: ...
+
+
+@overload
+def load_ckpt(path: Path, *, part: Literal["model"], model_template: PyTree) -> PyTree: ...
+
+
+@overload
+def load_ckpt(path: Path, *, part: Literal["opt"], optimizer_template: PyTree) -> optax.GradientTransformation: ...
+
+
+@overload
+def load_ckpt(path: Path, *, part: Literal["opt_state"], opt_state_template: PyTree) -> optax.OptState: ...
+
+
+@overload
+def load_ckpt(path: Path, *, part: Literal["state"]) -> State: ...
+
+
+@overload
+def load_ckpt(path: Path, *, part: Literal["config"]) -> DictConfig: ...
+
+
+def load_ckpt(
+    path: str | Path,
+    *,
+    part: CheckpointPart = "model",
+    model_template: PyTree | None = None,
+    optimizer_template: PyTree | None = None,
+    opt_state_template: PyTree | None = None,
+) -> (
+    tuple[PyTree, optax.GradientTransformation, optax.OptState, State, DictConfig]
+    | tuple[PyTree, State, DictConfig]
+    | PyTree
+    | optax.GradientTransformation
+    | optax.OptState
+    | State
+    | DictConfig
+):
+    with tarfile.open(path, "r:gz") as tar:
+
+        def get_model() -> PyTree:
+            if model_template is None:
+                raise ValueError("model_template must be provided to load model weights")
+            if (model := tar.extractfile("model")) is None:
+                raise ValueError(f"Checkpoint does not contain a model file: {path}")
+            return eqx.tree_deserialise_leaves(io.BytesIO(model.read()), model_template)
+
+        def get_opt() -> optax.GradientTransformation:
+            if optimizer_template is None:
+                raise ValueError("optimizer_template must be provided to load optimizer")
+            if (opt := tar.extractfile("optimizer")) is None:
+                raise ValueError(f"Checkpoint does not contain an optimizer file: {path}")
+            return eqx.tree_deserialise_leaves(io.BytesIO(opt.read()), optimizer_template)
+
+        def get_opt_state() -> optax.OptState:
+            if opt_state_template is None:
+                raise ValueError("opt_state_template must be provided to load optimizer state")
+            if (opt_state := tar.extractfile("opt_state")) is None:
+                raise ValueError(f"Checkpoint does not contain an optimizer state file: {path}")
+            return eqx.tree_deserialise_leaves(io.BytesIO(opt_state.read()), opt_state_template)
+
+        def get_state() -> State:
+            if (state := tar.extractfile("state")) is None:
+                raise ValueError(f"Checkpoint does not contain a state file: {path}")
+            return State.from_dict(**json.loads(state.read().decode()))
+
+        def get_config() -> DictConfig:
+            if (config := tar.extractfile("config")) is None:
+                raise ValueError(f"Checkpoint does not contain a config file: {path}")
+            return cast(DictConfig, OmegaConf.load(config))
+
+        match part:
+            case "model":
+                return get_model()
+            case "opt":
+                return get_opt()
+            case "opt_state":
+                return get_opt_state()
+            case "state":
+                return get_state()
+            case "config":
+                return get_config()
+            case "model_state_config":
+                return get_model(), get_state(), get_config()
+            case "all":
+                return get_model(), get_opt(), get_opt_state(), get_state(), get_config()
+            case _:
+                raise ValueError(f"Invalid checkpoint part: {part}")
+
+
 class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
     def __init__(self, config: Config) -> None:
         super().__init__(config)
@@ -82,149 +190,6 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
                 return True
         return False
 
-    @overload
-    def load_ckpt_with_template(
-        self,
-        path: Path,
-        *,
-        part: Literal["all"],
-        model_template: PyTree,
-        optimizer_template: PyTree,
-        opt_state_template: PyTree,
-    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State, Config]: ...
-
-    @overload
-    def load_ckpt_with_template(
-        self,
-        path: Path,
-        *,
-        part: Literal["model_state_config"],
-        model_template: PyTree,
-    ) -> tuple[PyTree, State, Config]: ...
-
-    @overload
-    def load_ckpt_with_template(
-        self,
-        path: Path,
-        *,
-        part: Literal["model"],
-        model_template: PyTree,
-    ) -> PyTree: ...
-
-    @overload
-    def load_ckpt_with_template(
-        self,
-        path: Path,
-        *,
-        part: Literal["opt"],
-        optimizer_template: PyTree,
-    ) -> optax.GradientTransformation: ...
-
-    @overload
-    def load_ckpt_with_template(
-        self,
-        path: Path,
-        *,
-        part: Literal["opt_state"],
-        opt_state_template: PyTree,
-    ) -> optax.OptState: ...
-
-    @overload
-    def load_ckpt_with_template(
-        self,
-        path: Path,
-        *,
-        part: Literal["state"],
-    ) -> State: ...
-
-    @overload
-    def load_ckpt_with_template(
-        self,
-        path: Path,
-        *,
-        part: Literal["config"],
-    ) -> Config: ...
-
-    def load_ckpt_with_template(
-        self,
-        path: Path,
-        *,
-        part: CheckpointPart = "all",
-        model_template: PyTree | None = None,
-        optimizer_template: PyTree | None = None,
-        opt_state_template: PyTree | None = None,
-    ) -> (
-        tuple[PyTree, optax.GradientTransformation, optax.OptState, State, Config]
-        | tuple[PyTree, State, Config]
-        | PyTree
-        | optax.GradientTransformation
-        | optax.OptState
-        | State
-        | Config
-    ):
-        """Load a checkpoint.
-
-        Args:
-            path: Path to the checkpoint directory
-            part: Which part of the checkpoint to load
-            model_template: Template model with correct structure but uninitialized weights
-            optimizer_template: Template optimizer with correct structure but uninitialized weights
-            opt_state_template: Template optimizer state with correct structure but uninitialized weights
-
-        Returns:
-            The requested checkpoint components
-        """
-        with tarfile.open(path, "r:gz") as tar:
-
-            def get_model() -> PyTree:
-                if model_template is None:
-                    raise ValueError("model_template must be provided to load model weights")
-                if (model := tar.extractfile("model")) is None:
-                    raise ValueError(f"Checkpoint does not contain a model file: {path}")
-                return eqx.tree_deserialise_leaves(io.BytesIO(model.read()), model_template)
-
-            def get_opt() -> optax.GradientTransformation:
-                if optimizer_template is None:
-                    raise ValueError("optimizer_template must be provided to load optimizer")
-                if (opt := tar.extractfile("optimizer")) is None:
-                    raise ValueError(f"Checkpoint does not contain an optimizer file: {path}")
-                return eqx.tree_deserialise_leaves(io.BytesIO(opt.read()), optimizer_template)
-
-            def get_opt_state() -> optax.OptState:
-                if opt_state_template is None:
-                    raise ValueError("opt_state_template must be provided to load optimizer state")
-                if (opt_state := tar.extractfile("opt_state")) is None:
-                    raise ValueError(f"Checkpoint does not contain an optimizer state file: {path}")
-                return eqx.tree_deserialise_leaves(io.BytesIO(opt_state.read()), opt_state_template)
-
-            def get_state() -> State:
-                if (state := tar.extractfile("state")) is None:
-                    raise ValueError(f"Checkpoint does not contain a state file: {path}")
-                return State.from_dict(**json.loads(state.read().decode()))
-
-            def get_config() -> Config:
-                if (config := tar.extractfile("config")) is None:
-                    raise ValueError(f"Checkpoint does not contain a config file: {path}")
-                return self.get_config(cast(DictConfig, OmegaConf.load(config)), use_cli=False)
-
-            match part:
-                case "model":
-                    return get_model()
-                case "opt":
-                    return get_opt()
-                case "opt_state":
-                    return get_opt_state()
-                case "state":
-                    return get_state()
-                case "config":
-                    return get_config()
-                case "model_state_config":
-                    return get_model(), get_state(), get_config()
-                case "all":
-                    return get_model(), get_opt(), get_opt_state(), get_state(), get_config()
-                case _:
-                    raise ValueError(f"Invalid checkpoint part: {part}")
-
     def save_checkpoint(
         self,
         model: PyTree | None = None,

xax/task/mixins/train.py

@@ -40,7 +40,7 @@ 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, CheckpointPart
+from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin, CheckpointPart, load_ckpt
 from xax.task.mixins.data_loader import DataloadersConfig, DataloadersMixin
 from xax.task.mixins.logger import LoggerConfig, LoggerMixin
 from xax.task.mixins.runnable import RunnableConfig, RunnableMixin
@@ -96,6 +96,12 @@ def batches_per_step_schedule(schedule: list[int] | None) -> list[int] | None:
     return list(itertools.accumulate([0] + schedule))
 
 
+def get_param_count(pytree: PyTree) -> int:
+    """Calculates the total number of parameters in a PyTree."""
+    leaves, _ = jax.tree.flatten(pytree)
+    return sum(x.size for x in leaves if isinstance(x, jnp.ndarray))
+
+
 class ValidStepTimer:
     def __init__(
         self,
@@ -360,6 +366,7 @@ class TrainMixin(
         model = self.get_model(key)
         state = State.init_state()
 
+        self.log_model_size(model)
         if not load_optimizer:
             return model, state
 
@@ -450,44 +457,43 @@ class TrainMixin(
         match part:
             case "model_state_config":
                 model_spec = eqx.filter_eval_shape(self.get_model, key)
-                return self.load_ckpt_with_template(path, part="model_state_config", model_template=model_spec)
+                model, state, config = load_ckpt(path, part="model_state_config", model_template=model_spec)
+                config = self.get_config(config, use_cli=False)
+                return model, state, config
 
             case "model":
                 model_spec = eqx.filter_eval_shape(self.get_model, key)
-                return self.load_ckpt_with_template(path, part="model", model_template=model_spec)
-
-            case "config":
-                return self.load_ckpt_with_template(path, part="config")
+                return load_ckpt(path, part="model", model_template=model_spec)
 
             case "opt":
                 optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                return self.load_ckpt_with_template(path, part="opt", optimizer_template=optimizer_spec)
+                return load_ckpt(path, part="opt", optimizer_template=optimizer_spec)
 
             case "opt_state":
                 if model is None:
                     model_spec = eqx.filter_eval_shape(self.get_model, key)
-                    model = self.load_ckpt_with_template(path, part="model", model_template=model_spec)
+                    model = load_ckpt(path, part="model", model_template=model_spec)
                 if optimizer is None:
                     optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                    optimizer = self.load_ckpt_with_template(path, part="opt", optimizer_template=optimizer_spec)
+                    optimizer = load_ckpt(path, part="opt", optimizer_template=optimizer_spec)
                 opt_state_spec = eqx.filter_eval_shape(self.get_initial_opt_state, model, optimizer)
-                return self.load_ckpt_with_template(path, part="opt_state", opt_state_template=opt_state_spec)
+                return load_ckpt(path, part="opt_state", opt_state_template=opt_state_spec)
 
             case "state":
-                return self.load_ckpt_with_template(path, part="state")
+                return load_ckpt(path, part="state")
 
             case "config":
-                return self.load_ckpt_with_template(path, part="config")
+                return self.get_config(load_ckpt(path, part="config"), use_cli=False)
 
             case "all":
                 model_spec = eqx.filter_eval_shape(self.get_model, key)
-                model = self.load_ckpt_with_template(path, part="model", model_template=model_spec)
+                model = load_ckpt(path, part="model", model_template=model_spec)
                 optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                optimizer = self.load_ckpt_with_template(path, part="opt", optimizer_template=optimizer_spec)
+                optimizer = load_ckpt(path, part="opt", optimizer_template=optimizer_spec)
                 opt_state_spec = eqx.filter_eval_shape(self.get_initial_opt_state, model, optimizer)
-                opt_state = self.load_ckpt_with_template(path, part="opt_state", opt_state_template=opt_state_spec)
-                state = self.load_ckpt_with_template(path, part="state")
-                config = self.load_ckpt_with_template(path, part="config")
+                opt_state = load_ckpt(path, part="opt_state", opt_state_template=opt_state_spec)
+                state = load_ckpt(path, part="state")
+                config = self.get_config(load_ckpt(path, part="config"), use_cli=False)
                 return model, optimizer, opt_state, state, config
 
             case _:
@@ -683,6 +689,9 @@ class TrainMixin(
         self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
         self.logger.log_file("info.json", get_info_json())
 
+    def log_model_size(self, model: PyTree) -> None:
+        logger.info("Model size: %s", f"{get_param_count(model):,}")
+
     def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401
         return eqx.is_inexact_array(item)
 

xax/utils/jaxpr.py

@@ -3,10 +3,10 @@
 from pathlib import Path
 
 import jax
-import jax.core
+import jax.extend.core
 
 
-def save_jaxpr_dot(closed_jaxpr: jax.core.ClosedJaxpr, filename: str | Path) -> None:
+def save_jaxpr_dot(closed_jaxpr: jax.extend.core.ClosedJaxpr, filename: str | Path) -> None:
     """Save the JAXPR to a DOT file.
 
     Example usage:
@@ -30,15 +30,15 @@ def save_jaxpr_dot(closed_jaxpr: jax.core.ClosedJaxpr, filename: str | Path) ->
     with open(filename, "w") as f:
         f.write("digraph Jaxpr {\n")
 
-        var_names: dict[jax.core.Var, str] = {}
+        var_names: dict[jax.extend.core.Var, str] = {}
         var_count = 0
 
-        def get_var_name(var: jax.core.Var) -> str:
+        def get_var_name(var: jax.extend.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):
+            if isinstance(var, jax.extend.core.Literal):
                 # Create a name based on the literal value
                 name = f"lit_{var.val}"
                 return name

xax/utils/pytree.py

@@ -57,7 +57,7 @@ def pytree_has_nans(pytree: PyTree) -> Array:
 
 def update_pytree(cond: Array, new: PyTree, original: PyTree) -> PyTree:
     """Update a pytree based on a condition."""
-    # Tricky, need use tree_map because where expects array leafs.
+    # Tricky, need use tree.map because where expects array leafs.
     return jax.tree.map(lambda x, y: jnp.where(cond, x, y), new, original)
 
 

xax/utils/types/frozen_dict.py

@@ -138,7 +138,7 @@ class FrozenDict(Mapping[K, V]):
 
 def unfreeze(x: FrozenDict[K, V] | dict[str, Any]) -> dict[Any, Any]:  # noqa: ANN401
     if isinstance(x, FrozenDict):
-        return jax.tree_util.tree_map(lambda y: y, x._dict)
+        return jax.tree.map(lambda y: y, x._dict)
     elif isinstance(x, dict):
         ys = {}
         for key, value in x.items():

{xax-0.2.6.dist-info → xax-0.2.7.dist-info}/METADATA

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

{xax-0.2.6.dist-info → xax-0.2.7.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=pvOmycOUli9h53XMTlinjJdY4zzw_fvWoF05UQwoItI,14225
+xax/__init__.py,sha256=94V0RHzCNC-aVXtpv5jtYaXJWytSJYOJCjUR69dIM1g,14428
 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
@@ -9,7 +9,7 @@ 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
 xax/nn/export.py,sha256=pRfM2B4hB2EvljysC6AjtgB_7Cn7JtaP3dhYU2stZtY,5545
-xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
+xax/nn/functions.py,sha256=Gig46ZjFxzXl3ImOFvpAO6oG6xrnAZywp1Ez8Gwy0BU,2711
 xax/nn/geom.py,sha256=rImNlkHWeoNcY7f84nknizJ6uzsrMhbAtKeb2xAWxNY,6215
 xax/nn/losses.py,sha256=Q_NVnm5n4UPBvp5nI_1aUptfXnqFYoUeFwySiyvopHg,272
 xax/nn/norm.py,sha256=WgZ3QCrUnf-YecwhEtVPcr99fKK3ECl_UeiAs2uv7oo,564
@@ -26,13 +26,13 @@ xax/task/launchers/cli.py,sha256=cK7Nm-3fO-W2gTxpn3FEThsT2NvneS2w0UjA1Nt-84A,140
 xax/task/launchers/single_process.py,sha256=IoML-30g5c526yxkpbWSOtG_KpNQMakT7xujzB1gIAo,846
 xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/loggers/callback.py,sha256=lyuZX6Bir7xJM07ifdQIl1jlclgkiS82UO9V4y7wgPs,1582
-xax/task/loggers/json.py,sha256=_tKum6jk_gqVzO-4MqSNXbE-Mmn-yJzkRAT-N1y2zes,4139
+xax/task/loggers/json.py,sha256=Ukbo6eAq9mSDmo7AqCVu9OFFjjSjcIsdeKP_1WQbyuw,4110
 xax/task/loggers/state.py,sha256=6bG-NRsSUzAukYiglCT0oDj8zRMpffH4e1TKWGw1x4k,959
 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=8Hi-2G0EA5OFRjgiOutlk7HgkD5b-0GHazOAYxnGytM,11409
+xax/task/mixins/checkpointing.py,sha256=zqospBFnTbGt_iriiduVfXazINPbzWpwmIs91KAniMY,10147
 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,25 +41,25 @@ 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=X3DIMPqpMRAJvFeeDaw5qFkH70R6yJelBAD8KHHwkUc,31196
+xax/task/mixins/train.py,sha256=t2KW18S9vpUuvmr3VAeKzSEcEdAD6fbi_fjk-KZ6ssA,31426
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 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/jaxpr.py,sha256=H7pWl48ROXIB1-ZPWYfOn-ou3EBMxYWIwc_A0reJQoo,2333
 xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
-xax/utils/pytree.py,sha256=VFWhT0MQ99KjQyEYM6NFbqYq4_hOZwB23uhowMB4U34,8754
+xax/utils/pytree.py,sha256=Q_5NCr90Wlqw1x4yK-FXfDVZEkNK2figmsUSZ4fnxYk,8754
 xax/utils/tensorboard.py,sha256=P0oIFvX2Qts1H4lkpizhRIpQdD0MNppVMeut0Z94yCs,19878
 xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,7066
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/utils/types/frozen_dict.py,sha256=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
+xax/utils/types/frozen_dict.py,sha256=s57XaTo2jgeT4_SzNQEjsYb4XrNZJwm1ca4ZdXSE5TY,4676
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.6.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.6.dist-info/METADATA,sha256=g--2dJ3PIYDOWURjNhVmkv5B_VUcI2xY0iBa4oGuuVc,1879
-xax-0.2.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.2.6.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.6.dist-info/RECORD,,
+xax-0.2.7.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.7.dist-info/METADATA,sha256=7otxcR5N4nVGyao5-NsrLfyuY1LqX8v2PWMEWB3sVy8,1879
+xax-0.2.7.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.2.7.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.7.dist-info/RECORD,,