xax 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.1.8"
+__version__ = "0.1.10"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -23,7 +23,6 @@ __all__ = [
     "get_run_dir",
     "load_user_config",
     "State",
-    "cast_phase",
     "FourierEmbeddings",
     "IdentityPositionalEmbeddings",
     "LearnedPositionalEmbeddings",
@@ -41,9 +40,6 @@ __all__ = [
     "load_eqx_mlp",
     "make_eqx_mlp",
     "save_eqx",
-    "export",
-    "export_flax",
-    "export_with_params",
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "quat_to_euler",
@@ -180,7 +176,6 @@ NAME_MAP: dict[str, str] = {
     "get_run_dir": "core.conf",
     "load_user_config": "core.conf",
     "State": "core.state",
-    "cast_phase": "core.state",
     "FourierEmbeddings": "nn.embeddings",
     "IdentityPositionalEmbeddings": "nn.embeddings",
     "LearnedPositionalEmbeddings": "nn.embeddings",
@@ -198,9 +193,6 @@ NAME_MAP: dict[str, str] = {
     "load_eqx_mlp": "nn.equinox",
     "make_eqx_mlp": "nn.equinox",
     "save_eqx": "nn.equinox",
-    "export": "nn.export",
-    "export_flax": "nn.export",
-    "export_with_params": "nn.export",
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "quat_to_euler": "nn.geom",
@@ -329,7 +321,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         get_run_dir,
         load_user_config,
     )
-    from xax.core.state import Phase, State, cast_phase
+    from xax.core.state import Phase, State
     from xax.nn.embeddings import (
         EmbeddingKind,
         FourierEmbeddings,
@@ -354,11 +346,6 @@ if IMPORT_ALL or TYPE_CHECKING:
         make_eqx_mlp,
         save_eqx,
     )
-    from xax.nn.export import (
-        export,
-        export_flax,
-        export_with_params,
-    )
     from xax.nn.geom import (
         euler_to_quat,
         get_projected_gravity_vector_from_quat,

xax/core/state.py

@@ -1,9 +1,10 @@
 """Defines a dataclass for keeping track of the current training state."""
 
 import time
-from dataclasses import dataclass
-from typing import Literal, NotRequired, TypedDict, cast, get_args
+from dataclasses import asdict, dataclass
+from typing import Any, Literal, NotRequired, TypedDict, Unpack, cast
 
+import jax
 from omegaconf import MISSING
 
 from xax.core.conf import field
@@ -11,12 +12,6 @@ from xax.core.conf import field
 Phase = Literal["train", "valid"]
 
 
-def cast_phase(raw_phase: str) -> Phase:
-    args = get_args(Phase)
-    assert raw_phase in args, f"Invalid phase: '{raw_phase}' Valid options are {args}"
-    return cast(Phase, raw_phase)
-
-
 class StateDict(TypedDict, total=False):
     num_steps: NotRequired[int]
     num_samples: NotRequired[int]
@@ -24,10 +19,11 @@ class StateDict(TypedDict, total=False):
     num_valid_samples: NotRequired[int]
     start_time_s: NotRequired[float]
     elapsed_time_s: NotRequired[float]
-    raw_phase: NotRequired[str]
+    phase: NotRequired[Phase]
 
 
-@dataclass
+@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")
@@ -35,15 +31,11 @@ class State:
     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")
-    raw_phase: str = field(MISSING, help="Current training phase")
+    _phase: int = field(MISSING, help="Current training phase")
 
     @property
     def phase(self) -> Phase:
-        return cast_phase(self.raw_phase)
-
-    @phase.setter
-    def phase(self, phase: Phase) -> None:
-        self.raw_phase = phase
+        return cast(Phase, ["train", "valid"][self._phase])
 
     @classmethod
     def init_state(cls) -> "State":
@@ -54,7 +46,7 @@ class State:
             num_valid_samples=0,
             start_time_s=time.time(),
             elapsed_time_s=0.0,
-            raw_phase="train",
+            _phase=0,
         )
 
     @property
@@ -69,3 +61,16 @@ class State:
                 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
+        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})

xax/task/base.py

@@ -16,7 +16,8 @@ from types import TracebackType
 from typing import Generic, Self, TypeVar, cast
 
 import jax
-from omegaconf import Container, DictConfig, OmegaConf
+from omegaconf import DictConfig, OmegaConf
+from omegaconf.base import SCMode
 
 from xax.core.state import State
 from xax.utils.text import camelcase_to_snakecase
@@ -66,9 +67,6 @@ class BaseTask(Generic[Config]):
 
         self.config = config
 
-        if isinstance(self.config, Container):
-            OmegaConf.resolve(self.config)
-
     def on_step_start(self, state: State) -> State:
         return state
 
@@ -195,7 +193,15 @@ class BaseTask(Generic[Config]):
                 cfg = OmegaConf.merge(cfg, *(get_config(path, task_path) for path in paths))
             cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(non_paths))
 
-        return cast(Config, cfg)
+        return cast(
+            Config,
+            OmegaConf.to_container(
+                cfg,
+                resolve=True,
+                throw_on_missing=True,
+                structured_config_mode=SCMode.INSTANTIATE,
+            ),
+        )
 
     @classmethod
     def config_str(cls, *cfgs: RawConfigType, use_cli: bool | list[str] = True) -> str:

xax/task/mixins/train.py

@@ -53,8 +53,10 @@ from xax.utils.experiments import (
     get_packages_with_versions,
     get_training_code,
 )
+from xax.utils.jax import jit as xax_jit
 from xax.utils.logging import LOG_STATUS
 from xax.utils.text import highlight_exception_message, show_info
+from xax.utils.types.frozen_dict import FrozenDict
 
 logger = logging.getLogger(__name__)
 
@@ -212,32 +214,31 @@ class TrainMixin(
 
     def on_step_end(self, state: State) -> State:
         state = super().on_step_end(state)
-        state.elapsed_time_s = time.time() - state.start_time_s
-        return state
+        return state.replace(elapsed_time_s=time.time() - state.start_time_s)
 
-    def log_train_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
+    def log_train_step(self, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State) -> None:
         """Override this function to do logging during the training phase.
 
         This function is called after the model forward pass and before the
         backward pass. It is called in the training phase.
 
         Args:
-            model: The current model.
             batch: The batch from the dataloader.
             output: The model output.
+            metrics: The metrics for the current batch.
             state: The current training state.
         """
 
-    def log_valid_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
+    def log_valid_step(self, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State) -> None:
         """Override this function to do logging during the validation phase.
 
         This function is called after the model forward pass. It is called in
         the validation phase.
 
         Args:
-            model: The current model.
             batch: The batch from the dataloader.
             output: The model output.
+            metrics: The metrics for the current batch.
             state: The current training state.
         """
 
@@ -248,18 +249,23 @@ class TrainMixin(
             for k, v in d.items():
                 self.logger.log_scalar(k, v, namespace=ns)
 
-    def log_step(self, model: PyTree, batch: Batch, output: Output, loss: Array, state: State) -> None:
+    def log_step(self, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State) -> None:
         phase = state.phase
 
-        self.logger.log_scalar("loss", loss, namespace="loss")
+        for k, v in metrics.items():
+            if v.size == 1:
+                self.logger.log_scalar(k, v.item())
+            else:
+                self.logger.log_histogram(k, v)
+
         self.log_state_timers(state)
 
         # Delegate to the appropriate logging function based on the phase.
         match phase:
             case "train":
-                self.log_train_step(model, batch, output, state)
+                self.log_train_step(batch, output, metrics, state)
             case "valid":
-                self.log_valid_step(model, batch, output, state)
+                self.log_valid_step(batch, output, metrics, state)
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
@@ -332,8 +338,7 @@ class TrainMixin(
 
         return model, optimizer, opt_state, state
 
-    @eqx.filter_jit
-    def get_output(self, model: PyTree, batch: Batch) -> Output:
+    def get_output(self, model: PyTree, batch: Batch, state: State) -> Output:
         """Gets the output from the model.
 
         By default, we assume the model is a function that takes the batch as
@@ -343,11 +348,11 @@ class TrainMixin(
         Args:
             model: The current model.
             batch: The current minibatch of samples.
+            state: The current training state.
         """
         raise NotImplementedError("`get_output` must be implemented by the subclass")
 
-    @eqx.filter_jit
-    def compute_loss(self, model: PyTree, batch: Batch, output: Output) -> Array:
+    def compute_loss(self, model: PyTree, batch: Batch, output: Output, state: State) -> Array:
         """Gets the loss for the current batch.
 
         By default, we assume the model is a function that takes the batch as
@@ -358,6 +363,7 @@ class TrainMixin(
             model: The current model.
             batch: The current minibatch of samples.
             output: The output from the model.
+            state: The current training state.
 
         Returns:
             The computed loss, as a tensor.
@@ -366,24 +372,59 @@ class TrainMixin(
             raise ValueError(f"When model output is not the loss, you must override `compute_loss`. Got {type(output)}")
         return output
 
-    @eqx.filter_jit
-    def get_output_and_loss(self, model: PyTree, batch: Batch) -> tuple[Array, Output]:
-        output = self.get_output(model, batch)
-        loss = self.compute_loss(model, batch, output)
-        return loss, output
+    def compute_metrics(
+        self,
+        model: PyTree,
+        batch: Batch,
+        output: Output,
+        loss: Array,
+        state: State,
+    ) -> dict[str, Array]:
+        """Computes the metrics for the current batch.
+
+        Args:
+            model: The current model.
+            batch: The current minibatch of samples.
+            output: The output from the model.
+            loss: The loss for the current batch.
+            state: The current training state.
+
+        Returns:
+            A dictionary of metrics.
+        """
+        return {
+            "loss": loss,
+        }
+
+    @xax_jit(static_argnames=["self", "model_static"])
+    def get_output_and_loss(
+        self,
+        model_arr: PyTree,
+        model_static: PyTree,
+        batch: Batch,
+        state: State,
+    ) -> tuple[Array, tuple[Output, FrozenDict[str, Array]]]:
+        model = eqx.combine(model_arr, model_static)
+        output = self.get_output(model, batch, state)
+        loss = self.compute_loss(model, batch, output, state)
+        metrics = self.compute_metrics(model, batch, output, loss, state)
+        return loss, (output, FrozenDict(metrics))
 
-    @eqx.filter_jit
     def update(
         self,
-        model: PyTree,
+        model_arr: PyTree,
+        model_static: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
-    ) -> tuple[Array, PyTree, optax.OptState, Output]:
-        (loss, output), grads = eqx.filter_value_and_grad(self.get_output_and_loss, has_aux=True)(model, batch)
-        updates, opt_state = optimizer.update(grads, opt_state, model)
-        model = eqx.apply_updates(model, updates)
-        return loss, model, opt_state, output
+        state: State,
+    ) -> tuple[PyTree, optax.OptState, Output, FrozenDict[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)
+        grads, (output, metrics) = grad_fn(model_arr, model_static, batch, state)
+        updates, opt_state = optimizer.update(grads, opt_state, model_arr)
+        model_arr = eqx.apply_updates(model_arr, updates)
+        return model_arr, opt_state, output, metrics
 
     def get_size_of_batch(self, batch: Batch) -> int | None:
         """Gets the batch size for the current batch.
@@ -457,21 +498,32 @@ class TrainMixin(
         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))
 
-    @eqx.filter_jit
+    def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401
+        return eqx.is_inexact_array(item)
+
+    @xax_jit(static_argnames=["self", "model_static", "optimizer"])
     def train_step(
         self,
-        model: PyTree,
+        model_arr: PyTree,
+        model_static: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
-    ) -> tuple[PyTree, optax.OptState, Array, Output]:
-        loss, model, opt_state, output = self.update(model, optimizer, opt_state, batch)
-        return model, opt_state, loss, output
+        state: State,
+    ) -> tuple[PyTree, optax.OptState, Output, FrozenDict[str, Array]]:
+        model_arr, opt_state, output, metrics = self.update(model_arr, model_static, optimizer, opt_state, batch, state)
+        return model_arr, opt_state, output, metrics
 
-    @eqx.filter_jit
-    def val_step(self, model: PyTree, batch: Batch) -> tuple[PyTree, Array, Output]:
-        loss, output = eqx.filter_jit(self.get_output_and_loss)(model, batch)
-        return model, loss, output
+    @xax_jit(static_argnames=["self", "model_static"])
+    def val_step(
+        self,
+        model_arr: PyTree,
+        model_static: PyTree,
+        batch: Batch,
+        state: State,
+    ) -> tuple[Output, FrozenDict[str, Array]]:
+        _, (output, metrics) = self.get_output_and_loss(model_arr, model_static, batch, state)
+        return output, metrics
 
     def train_loop(
         self,
@@ -482,36 +534,46 @@ class TrainMixin(
         valid_pf: Iterator[Batch],
         state: State,
     ) -> None:
+        model_arr, model_static = eqx.partition(model, self.model_partition_fn)
+
         while not self.is_training_over(state):
             if self.valid_step_timer.is_valid_step(state):
                 valid_batch = next(valid_pf)
-                with self.step_context("model_step"):
-                    model, loss, output = self.val_step(model, valid_batch)
+                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),
+                )
 
-                # Perform logging.
-                with self.step_context("write_logs"):
-                    state.phase = "valid"
-                    self.log_step(model, valid_batch, output, loss, state)
-                    state.num_valid_samples += 1
+                output, metrics = self.val_step(model_arr, model_static, valid_batch, state)
+                self.log_step(valid_batch, output, metrics, state)
 
             state = self.on_step_start(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)
-
-            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
+            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),
+            )
+
+            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(train_batch, output, metrics, state)
 
             state = self.on_step_end(state)
 
             if self.should_checkpoint(state):
+                model = eqx.combine(model_arr, model_static)
                 self.save_checkpoint(model, optimizer, opt_state, state)
 
         # After finishing training, save the final checkpoint.
+        model = eqx.combine(model_arr, model_static)
         self.save_checkpoint(model, optimizer, opt_state, state)
 
     @contextlib.contextmanager

xax/task/script.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from typing import Generic, TypeVar
 
+import jax
+
 from xax.task.base import BaseConfig, BaseTask
 from xax.task.mixins import (
     ArtifactsConfig,
@@ -20,6 +22,7 @@ from xax.task.mixins import (
 )
 
 
+@jax.tree_util.register_dataclass
 @dataclass(kw_only=True)
 class ScriptConfig(
     CPUStatsConfig,

xax/utils/pytree.py

@@ -31,7 +31,7 @@ def slice_array(x: Array, start: Array, slice_length: int) -> Array:
 
 def slice_pytree(pytree: PyTree, start: Array, slice_length: int) -> PyTree:
     """Get a slice of a pytree."""
-    return jax.tree_util.tree_map(lambda x: slice_array(x, start, slice_length), pytree)
+    return jax.tree.map(lambda x: slice_array(x, start, slice_length), pytree)
 
 
 def flatten_array(x: Array, flatten_size: int) -> Array:
@@ -43,14 +43,14 @@ def flatten_array(x: Array, flatten_size: int) -> Array:
 
 def flatten_pytree(pytree: PyTree, flatten_size: int) -> PyTree:
     """Flatten a pytree into a (flatten_size, ...) pytree."""
-    return jax.tree_util.tree_map(lambda x: flatten_array(x, flatten_size), pytree)
+    return jax.tree.map(lambda x: flatten_array(x, flatten_size), pytree)
 
 
 def pytree_has_nans(pytree: PyTree) -> Array:
     """Check if a pytree has any NaNs."""
     has_nans = jax.tree_util.tree_reduce(
         lambda a, b: jnp.logical_or(a, b),
-        jax.tree_util.tree_map(lambda x: jnp.any(jnp.isnan(x)), pytree),
+        jax.tree.map(lambda x: jnp.any(jnp.isnan(x)), pytree),
     )
     return has_nans
 
@@ -58,13 +58,13 @@ 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.
-    return jax.tree_util.tree_map(lambda x, y: jnp.where(cond, x, y), new, original)
+    return jax.tree.map(lambda x, y: jnp.where(cond, x, y), new, original)
 
 
 def compute_nan_ratio(pytree: PyTree) -> Array:
     """Computes the ratio of NaNs vs non-NaNs in a given PyTree."""
-    nan_counts = jax.tree_util.tree_map(lambda x: jnp.sum(jnp.isnan(x)), pytree)
-    total_counts = jax.tree_util.tree_map(lambda x: x.size, pytree)
+    nan_counts = jax.tree.map(lambda x: jnp.sum(jnp.isnan(x)), pytree)
+    total_counts = jax.tree.map(lambda x: x.size, pytree)
 
     total_nans = jax.tree_util.tree_reduce(lambda a, b: a + b, nan_counts, 0)
     total_elements = jax.tree_util.tree_reduce(lambda a, b: a + b, total_counts, 0)
@@ -118,7 +118,7 @@ def reshuffle_pytree(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArr
         # Reshape back to the original shape
         return permuted.reshape(orig_shape)
 
-    return jax.tree_util.tree_map(permute_array, data)
+    return jax.tree.map(permute_array, data)
 
 
 def reshuffle_pytree_independently(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArray) -> PyTree:
@@ -133,7 +133,7 @@ def reshuffle_pytree_independently(data: PyTree, batch_shape: tuple[int, ...], r
             return x[tuple(idx_grids)]
         return x
 
-    return jax.tree_util.tree_map(permute_array, data)
+    return jax.tree.map(permute_array, data)
 
 
 TransposeResult = tuple[PyTree, tuple[int, ...], tuple[int, ...]]
@@ -215,7 +215,7 @@ def reshuffle_pytree_along_dims(
                 transpose_info[path] = (transpose_order, original_shape)
         return x
 
-    jax.tree_util.tree_map_with_path(prepare_for_shuffle, data)
+    jax.tree.map_with_path(prepare_for_shuffle, data)
 
     # Create a transposed pytree
     def get_transposed(path: PathType, x: PyTree) -> PyTree:
@@ -223,7 +223,7 @@ def reshuffle_pytree_along_dims(
             return transposed_data[path]
         return x
 
-    transposed_pytree = jax.tree_util.tree_map_with_path(get_transposed, data)
+    transposed_pytree = jax.tree.map_with_path(get_transposed, data)
 
     # Reshuffle the transposed pytree along the leading dimensions
     reshuffled_transposed = reshuffle_pytree(transposed_pytree, shape_dims, rng)
@@ -235,4 +235,4 @@ def reshuffle_pytree_along_dims(
             return transpose_back(x, transpose_order, original_shape)
         return x
 
-    return jax.tree_util.tree_map_with_path(restore_transpose, reshuffled_transposed)
+    return jax.tree.map_with_path(restore_transpose, reshuffled_transposed)

{xax-0.1.8.dist-info → xax-0.1.10.dist-info}/METADATA

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

{xax-0.1.8.dist-info → xax-0.1.10.dist-info}/RECORD

@@ -1,10 +1,10 @@
-xax/__init__.py,sha256=9WNjoeAF7enu7YXQqshpVG1FucGdSkxwrRa-ELDDuUs,13713
+xax/__init__.py,sha256=bvOBMlEVA46I7ILGfk5AbpwpcdTAjw-4vWI7ci7L7-g,13392
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
-xax/core/state.py,sha256=y123fL7pMgk25TPG6KN0LRIF_eYnD9eP7OfqtoQJGNE,2178
+xax/core/state.py,sha256=WwW0qDm-be9MMOT-bGWEFvaWF4iq2FP9xRSn1zq_4A8,2507
 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
@@ -14,9 +14,9 @@ xax/nn/geom.py,sha256=eK7I8fUHBc3FT7zpm5Yf__bXFQ4LtX6sa17-DxojLTo,3202
 xax/nn/norm.py,sha256=cDmYf5CtyzmuCiWdSP5nr8nZKQOmaZueDQXMPnThg6c,548
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/base.py,sha256=4fUjrG-llQpeESQuaQbww4M6WR6djjTK89fY20UV9zU,7610
+xax/task/base.py,sha256=E4l1yCrAkM2TVTbVYrmk6BoVHMkbD4IYsTT921XOyi0,7760
 xax/task/logger.py,sha256=1SZjVC6UCtZUoMPcpp3ckotL324QDeYDvHVhf5MHVqg,36271
-xax/task/script.py,sha256=zt36Sobdoer86gXHqc4sMAW7bqZRVl6IEExuQZH2USk,926
+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
 xax/task/launchers/base.py,sha256=8LB_r6YISKu1vq1zk3aVYmiedRr9MxE5IMRocs6unFI,731
@@ -39,7 +39,7 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=d1opVgvc6bOFXb7R58b07F4P5lbSZIzYaajtE0eBbpw,1477
 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=vsH_QpyrThlh9AzWnyvDJv58Y8U_516oi8gmMq_0iMg,22333
+xax/task/mixins/train.py,sha256=jAzc9RD25DbhekvItzsRQQrK9aEwtA_sXy0m2Hfkuxo,24594
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=9WlCrEqbq-SVXPEM4rhsLYERH97XNX7XSYLSI3sgKGk,1619
 xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
@@ -48,7 +48,7 @@ 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
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
-xax/utils/pytree.py,sha256=7GjQoPc_ZSZt3QS_9qXoBWl1jfMp1qZa7aViQoWJ0OQ,8864
+xax/utils/pytree.py,sha256=VFWhT0MQ99KjQyEYM6NFbqYq4_hOZwB23uhowMB4U34,8754
 xax/utils/tensorboard.py,sha256=21czW8WC2SAmwEhz6RLJc_q5HFvNKM4iR1ZycSO5qPE,17058
 xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -56,8 +56,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.1.8.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.8.dist-info/METADATA,sha256=wnBSNRByXJzgQPuZqNWooidfFdqcT4w8gbwlBgzbJk8,1877
-xax-0.1.8.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.8.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.8.dist-info/RECORD,,
+xax-0.1.10.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.1.10.dist-info/METADATA,sha256=kJ1lxZ6cWrtJ5R-adTorzEE_1l0VRJ67xfuBjYXG9Vo,1878
+xax-0.1.10.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.1.10.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.1.10.dist-info/RECORD,,