xax 0.2.13__py3-none-any.whl → 0.2.14__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.13"
+__version__ = "0.2.14"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

xax/task/mixins/checkpointing.py

@@ -6,7 +6,7 @@ import logging
 import tarfile
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Generic, Literal, TypeVar, cast, overload
+from typing import Generic, Literal, Sequence, TypeVar, cast, overload
 
 import equinox as eqx
 import jax
@@ -57,10 +57,10 @@ 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]: ...
+    model_templates: Sequence[PyTree],
+    optimizer_templates: Sequence[optax.GradientTransformation],
+    opt_state_templates: Sequence[optax.OptState],
+) -> tuple[list[PyTree], list[optax.GradientTransformation], list[optax.OptState], State, DictConfig]: ...
 
 
 @overload
@@ -68,20 +68,35 @@ def load_ckpt(
     path: Path,
     *,
     part: Literal["model_state_config"],
-    model_template: PyTree,
-) -> tuple[PyTree, State, DictConfig]: ...
+    model_templates: Sequence[PyTree],
+) -> tuple[list[PyTree], State, DictConfig]: ...
 
 
 @overload
-def load_ckpt(path: Path, *, part: Literal["model"], model_template: PyTree) -> PyTree: ...
+def load_ckpt(
+    path: Path,
+    *,
+    part: Literal["model"],
+    model_templates: Sequence[PyTree],
+) -> list[PyTree]: ...
 
 
 @overload
-def load_ckpt(path: Path, *, part: Literal["opt"], optimizer_template: PyTree) -> optax.GradientTransformation: ...
+def load_ckpt(
+    path: Path,
+    *,
+    part: Literal["opt"],
+    optimizer_templates: Sequence[optax.GradientTransformation],
+) -> list[optax.GradientTransformation]: ...
 
 
 @overload
-def load_ckpt(path: Path, *, part: Literal["opt_state"], opt_state_template: PyTree) -> optax.OptState: ...
+def load_ckpt(
+    path: Path,
+    *,
+    part: Literal["opt_state"],
+    opt_state_templates: Sequence[optax.OptState],
+) -> list[optax.OptState]: ...
 
 
 @overload
@@ -96,40 +111,49 @@ 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,
+    model_templates: Sequence[PyTree] | None = None,
+    optimizer_templates: Sequence[optax.GradientTransformation] | None = None,
+    opt_state_templates: Sequence[optax.OptState] | None = None,
 ) -> (
-    tuple[PyTree, optax.GradientTransformation, optax.OptState, State, DictConfig]
-    | tuple[PyTree, State, DictConfig]
-    | PyTree
-    | optax.GradientTransformation
-    | optax.OptState
+    tuple[list[PyTree], list[optax.GradientTransformation], list[optax.OptState], State, DictConfig]
+    | tuple[list[PyTree], State, DictConfig]
+    | list[PyTree]
+    | list[optax.GradientTransformation]
+    | list[optax.OptState]
     | State
     | DictConfig
 ):
     with tarfile.open(path, "r:gz") as tar:
 
-        def get_model() -> PyTree:
-            if model_template is None:
+        def get_model() -> list[PyTree]:
+            if model_templates 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:
+            models: list[PyTree] = []
+            for i, model_template in enumerate(model_templates):
+                if (model := tar.extractfile(f"model_{i}")) is None:
+                    raise ValueError(f"Checkpoint does not contain a model file: {path}")
+                models.append(eqx.tree_deserialise_leaves(io.BytesIO(model.read()), model_template))
+            return models
+
+        def get_opt() -> list[optax.GradientTransformation]:
+            if optimizer_templates 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:
+            opts: list[optax.GradientTransformation] = []
+            for i, optimizer_template in enumerate(optimizer_templates):
+                if (opt := tar.extractfile(f"optimizer_{i}")) is None:
+                    raise ValueError(f"Checkpoint does not contain an optimizer file: {path}")
+                opts.append(eqx.tree_deserialise_leaves(io.BytesIO(opt.read()), optimizer_template))
+            return opts
+
+        def get_opt_state() -> list[optax.OptState]:
+            if opt_state_templates 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)
+            opt_states: list[optax.OptState] = []
+            for i, opt_state_template in enumerate(opt_state_templates):
+                if (opt_state := tar.extractfile(f"opt_state_{i}")) is None:
+                    raise ValueError(f"Checkpoint does not contain an optimizer state file: {path}")
+                opt_states.append(eqx.tree_deserialise_leaves(io.BytesIO(opt_state.read()), opt_state_template))
+            return opt_states
 
         def get_state() -> State:
             if (state := tar.extractfile("state")) is None:
@@ -192,20 +216,20 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
 
     def save_checkpoint(
         self,
-        model: PyTree | None = None,
-        optimizer: optax.GradientTransformation | None = None,
-        opt_state: optax.OptState | None = None,
+        models: Sequence[PyTree] | None = None,
+        optimizers: Sequence[optax.GradientTransformation] | None = None,
+        opt_states: Sequence[optax.OptState] | None = None,
         aux_data: PyTree | None = None,
         state: State | None = None,
     ) -> Path:
         """Save a checkpoint.
 
         Args:
-            model: The model to save
-            state: The current training state
-            optimizer: The optimizer to save
+            models: The models to save
+            optimizers: The optimizers to save
+            opt_states: The optimizer states to save
             aux_data: Additional data to save
-            opt_state: The optimizer state to save
+            state: The current training state
 
         Returns:
             Path to the saved checkpoint
@@ -235,22 +259,25 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
                 tar.addfile(tarinfo, buf)
 
             # Save model using Equinox
-            if model is not None:
-                with io.BytesIO() as buf:
-                    eqx.tree_serialise_leaves(buf, model)
-                    add_file("model", buf)
+            if models is not None:
+                for i, model in enumerate(models):
+                    with io.BytesIO() as buf:
+                        eqx.tree_serialise_leaves(buf, model)
+                        add_file(f"model_{i}", buf)
 
             # Save optimizer using Equinox
-            if optimizer is not None:
-                with io.BytesIO() as buf:
-                    eqx.tree_serialise_leaves(buf, optimizer)
-                    add_file("optimizer", buf)
+            if optimizers is not None:
+                for i, optimizer in enumerate(optimizers):
+                    with io.BytesIO() as buf:
+                        eqx.tree_serialise_leaves(buf, optimizer)
+                        add_file(f"optimizer_{i}", buf)
 
             # Save optimizer state using Equinox
-            if opt_state is not None:
-                with io.BytesIO() as buf:
-                    eqx.tree_serialise_leaves(buf, opt_state)
-                    add_file("opt_state", buf)
+            if opt_states is not None:
+                for i, opt_state in enumerate(opt_states):
+                    with io.BytesIO() as buf:
+                        eqx.tree_serialise_leaves(buf, opt_state)
+                        add_file(f"opt_state_{i}", buf)
 
             # Save aux data using Equinox.
             if aux_data is not None:

xax/task/mixins/train.py

@@ -310,23 +310,46 @@ class TrainMixin(
         self.write_logs(state)
 
     @abstractmethod
-    def get_model(self, key: PRNGKeyArray) -> PyTree:
+    def get_model(self, key: PRNGKeyArray) -> PyTree | Sequence[PyTree]:
         """Returns the Equinox model to train.
 
         Returns:
             The model to train.
         """
 
+    def _get_models(self, key: PRNGKeyArray) -> list[PyTree]:
+        models = self.get_model(key)
+        if isinstance(models, Sequence):
+            models = list(models)
+        elif isinstance(models, eqx.Module):
+            models = [models]
+        else:
+            logger.warning("Model is not a sequence or an eqx.Module, wrapping it in a list anyway")
+            models = [models]
+        return models
+
     @abstractmethod
-    def get_optimizer(self) -> optax.GradientTransformation:
+    def get_optimizer(self) -> optax.GradientTransformation | Sequence[optax.GradientTransformation]:
         """Gets the optimizer for the model.
 
         Returns:
             The optimizer to use to train the model.
         """
 
-    def get_initial_opt_state(self, model: PyTree, optimizer: optax.GradientTransformation) -> optax.OptState:
-        return optimizer.init(eqx.filter(model, eqx.is_array))
+    def _get_optimizers(self) -> list[optax.GradientTransformation]:
+        optimizers = self.get_optimizer()
+        if isinstance(optimizers, optax.GradientTransformation):
+            optimizers = [optimizers]
+        elif isinstance(optimizers, Sequence):
+            optimizers = list(optimizers)
+        return optimizers
+
+    def get_initial_opt_state(
+        self,
+        models: list[PyTree],
+        optimizers: list[optax.GradientTransformation],
+    ) -> list[optax.OptState]:
+        return [opt.init(eqx.filter(model, eqx.is_array)) for model, opt in zip(models, optimizers, strict=True)]
 
     @overload
     def load_initial_state(
@@ -346,7 +369,10 @@ class TrainMixin(
         self,
         key: PRNGKeyArray,
         load_optimizer: bool = False,
-    ) -> tuple[PyTree, State] | tuple[PyTree, optax.GradientTransformation, optax.OptState, State]:
+    ) -> (
+        tuple[list[PyTree], State]
+        | tuple[list[PyTree], list[optax.GradientTransformation], list[optax.OptState], State]
+    ):
         init_ckpt_path = self.get_init_ckpt_path()
 
         if init_ckpt_path is not None:
@@ -364,16 +390,17 @@ class TrainMixin(
             return model, optimizer, opt_state, state
 
         logger.info("Starting a new training run")
-        model = self.get_model(key)
+        models = self._get_models(key)
         state = State.init_state()
 
         if not load_optimizer:
-            return model, state
+            return models, state
 
-        optimizer = self.get_optimizer()
-        opt_state = self.get_initial_opt_state(model, optimizer)
+        # Gets the optimizer(s) for the model.
+        optimizers = self._get_optimizers()
+        opt_states = self.get_initial_opt_state(models, optimizers)
 
-        return model, optimizer, opt_state, state
+        return models, optimizers, opt_states, state
 
     @overload
     def load_ckpt(
@@ -381,7 +408,7 @@ class TrainMixin(
         path: Path,
         *,
         part: Literal["all"],
-    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State, Config]: ...
+    ) -> tuple[list[PyTree], list[optax.GradientTransformation], list[optax.OptState], State, Config]: ...
 
     @overload
     def load_ckpt(
@@ -389,7 +416,7 @@ class TrainMixin(
         path: Path,
         *,
         part: Literal["model_state_config"],
-    ) -> tuple[PyTree, State, Config]: ...
+    ) -> tuple[list[PyTree], State, Config]: ...
 
     @overload
     def load_ckpt(
@@ -397,7 +424,7 @@ class TrainMixin(
         path: Path,
         *,
         part: Literal["model"],
-    ) -> PyTree: ...
+    ) -> list[PyTree]: ...
 
     @overload
     def load_ckpt(
@@ -405,7 +432,7 @@ class TrainMixin(
         path: Path,
         *,
         part: Literal["opt"],
-    ) -> optax.GradientTransformation: ...
+    ) -> list[optax.GradientTransformation]: ...
 
     @overload
     def load_ckpt(
@@ -415,7 +442,7 @@ class TrainMixin(
         part: Literal["opt_state"],
         model: PyTree | None = None,
         optimizer: optax.GradientTransformation | None = None,
-    ) -> optax.OptState: ...
+    ) -> list[optax.OptState]: ...
 
     @overload
     def load_ckpt(
@@ -423,7 +450,7 @@ class TrainMixin(
         path: Path,
         *,
         part: Literal["state"],
-    ) -> State: ...
+    ) -> list[State]: ...
 
     @overload
     def load_ckpt(
@@ -431,7 +458,7 @@ class TrainMixin(
         path: Path,
         *,
         part: Literal["config"],
-    ) -> Config: ...
+    ) -> list[Config]: ...
 
     def load_ckpt(
         self,
@@ -441,11 +468,11 @@ class TrainMixin(
         model: PyTree | None = None,
         optimizer: optax.GradientTransformation | None = None,
     ) -> (
-        tuple[PyTree, optax.GradientTransformation, optax.OptState, State, Config]
-        | tuple[PyTree, State, Config]
-        | PyTree
-        | optax.GradientTransformation
-        | optax.OptState
+        tuple[list[PyTree], list[optax.GradientTransformation], list[optax.OptState], State, Config]
+        | tuple[list[PyTree], State, Config]
+        | list[PyTree]
+        | list[optax.GradientTransformation]
+        | list[optax.OptState]
         | State
         | Config
     ):
@@ -456,28 +483,28 @@ class TrainMixin(
 
         match part:
             case "model_state_config":
-                model_spec = eqx.filter_eval_shape(self.get_model, key)
-                model, state, config = load_ckpt(path, part="model_state_config", model_template=model_spec)
+                model_specs = eqx.filter_eval_shape(self._get_models, key)
+                model, state, config = load_ckpt(path, part="model_state_config", model_templates=model_specs)
                 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 load_ckpt(path, part="model", model_template=model_spec)
+                model_specs = eqx.filter_eval_shape(self._get_models, key)
+                return load_ckpt(path, part="model", model_templates=model_specs)
 
             case "opt":
-                optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                return load_ckpt(path, part="opt", optimizer_template=optimizer_spec)
+                optimizer_specs = eqx.filter_eval_shape(self._get_optimizers)
+                return load_ckpt(path, part="opt", optimizer_templates=optimizer_specs)
 
             case "opt_state":
                 if model is None:
-                    model_spec = eqx.filter_eval_shape(self.get_model, key)
-                    model = load_ckpt(path, part="model", model_template=model_spec)
+                    model_specs = eqx.filter_eval_shape(self._get_models, key)
+                    model = load_ckpt(path, part="model", model_templates=model_specs)
                 if optimizer is None:
-                    optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                    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 load_ckpt(path, part="opt_state", opt_state_template=opt_state_spec)
+                    optimizer_specs = eqx.filter_eval_shape(self._get_optimizers)
+                    optimizer = load_ckpt(path, part="opt", optimizer_templates=optimizer_specs)
+                opt_state_specs = eqx.filter_eval_shape(self.get_initial_opt_state, model, optimizer)
+                return load_ckpt(path, part="opt_state", opt_state_templates=opt_state_specs)
 
             case "state":
                 return load_ckpt(path, part="state")
@@ -486,12 +513,12 @@ class TrainMixin(
                 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 = load_ckpt(path, part="model", model_template=model_spec)
-                optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                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 = load_ckpt(path, part="opt_state", opt_state_template=opt_state_spec)
+                model_specs = eqx.filter_eval_shape(self._get_models, key)
+                model = load_ckpt(path, part="model", model_templates=model_specs)
+                optimizer_specs = eqx.filter_eval_shape(self._get_optimizers)
+                optimizer = load_ckpt(path, part="opt", optimizer_templates=optimizer_specs)
+                opt_state_specs = eqx.filter_eval_shape(self.get_initial_opt_state, model, optimizer)
+                opt_state = load_ckpt(path, part="opt_state", opt_state_templates=opt_state_specs)
                 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
@@ -718,14 +745,22 @@ class TrainMixin(
 
     def train_loop(
         self,
-        model: PyTree,
-        optimizer: optax.GradientTransformation,
-        opt_state: optax.OptState,
+        models: Sequence[PyTree],
+        optimizers: Sequence[optax.GradientTransformation],
+        opt_states: Sequence[optax.OptState],
         train_pf: Iterator[Batch],
         valid_pf: Iterator[Batch],
         state: State,
     ) -> None:
-        model_arr, model_static = eqx.partition(model, self.model_partition_fn)
+        if len(models) != 1 or len(optimizers) != 1 or len(opt_states) != 1:
+            raise ValueError(
+                "Vanilla training expects a single model, optimizer and optimizer state. "
+                f"Found {len(models)} models, {len(optimizers)} optimizers and {len(opt_states)} optimizer states."
+            )
+
+        model_arr, model_static = eqx.partition(models[0], self.model_partition_fn)
+        optimizer = optimizers[0]
+        opt_state = opt_states[0]
 
         while not self.is_training_over(state):
             valid_step = self.valid_step_timer(state)
@@ -773,11 +808,11 @@ class TrainMixin(
 
             if self.should_checkpoint(state):
                 model = eqx.combine(model_arr, model_static)
-                self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
+                self.save_checkpoint(models=[model], optimizers=[optimizer], opt_states=[opt_state], state=state)
 
         # After finishing training, save the final checkpoint.
         model = eqx.combine(model_arr, model_static)
-        self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
+        self.save_checkpoint(models=[model], optimizers=[optimizer], opt_states=[opt_state], state=state)
 
     @contextlib.contextmanager
     def get_train_iterator(self, key: PRNGKeyArray) -> Generator[Iterator[Batch], None, None]:
@@ -841,14 +876,14 @@ 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, load_optimizer=True)
-            logger.info("Model size: %s", f"{get_pytree_param_count(model):,}")
-            logger.info("Optimizer size: %s", f"{get_pytree_param_count(optimizer):,}")
+            models, optimizers, opt_states, state = self.load_initial_state(model_key, load_optimizer=True)
+            logger.info("Model size: %s", f"{get_pytree_param_count(models):,}")
+            logger.info("Optimizer size: %s", f"{get_pytree_param_count(optimizers):,}")
 
             state = self.on_training_start(state)
 
             def on_exit() -> None:
-                self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
+                self.save_checkpoint(models=models, optimizers=optimizers, opt_states=opt_states, state=state)
 
             # Handle user-defined interrupts during the training loop.
             self.add_signal_handler(on_exit, signal.SIGUSR1, signal.SIGTERM)
@@ -857,9 +892,9 @@ class TrainMixin(
             with self.get_train_iterator(tkey) as train_pf, self.get_valid_iterator(vkey) as valid_pf:
                 try:
                     self.train_loop(
-                        model=model,
-                        optimizer=optimizer,
-                        opt_state=opt_state,
+                        models=models,
+                        optimizers=optimizers,
+                        opt_states=opt_states,
                         train_pf=train_pf,
                         valid_pf=valid_pf,
                         state=state,
@@ -869,7 +904,7 @@ class TrainMixin(
                     if is_master():
                         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)
+                    self.save_checkpoint(models=models, optimizers=optimizers, opt_states=opt_states, state=state)
 
                 except (KeyboardInterrupt, bdb.BdbQuit):
                     if is_master():
@@ -879,7 +914,7 @@ class TrainMixin(
                     exception_tb = textwrap.indent(highlight_exception_message(traceback.format_exc()), "  ")
                     sys.stdout.write(f"Caught exception during training loop:\n\n{exception_tb}\n")
                     sys.stdout.flush()
-                    self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
+                    self.save_checkpoint(models=models, optimizers=optimizers, opt_states=opt_states, state=state)
 
                 finally:
                     state = self.on_training_end(state)

{xax-0.2.13.dist-info → xax-0.2.14.dist-info}/METADATA

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

{xax-0.2.13.dist-info → xax-0.2.14.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=33wIwGeXDFReg2ZnFqUHfSybj5cKyMqnI8ncj8-9yVg,15510
+xax/__init__.py,sha256=9RtfWhU2Qb-YlZurazJ_GKpSpPaGs0nwu9-nmzbC0Pk,15510
 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
@@ -32,7 +32,7 @@ xax/task/loggers/stdout.py,sha256=giKSW2R83YkgRefm3BLkE7t8Pbj5Dux4AgsdJxYIbGo,66
 xax/task/loggers/tensorboard.py,sha256=sRyBbeBeVXDTYhPZIKIapW0JEfL9hqqzhNTeIcSd374,8883
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=Ma7fwsp-SA1w6GcuBSskszj5TB83yxYJm4Ns_EnqkI4,3018
-xax/task/mixins/checkpointing.py,sha256=zqospBFnTbGt_iriiduVfXazINPbzWpwmIs91KAniMY,10147
+xax/task/mixins/checkpointing.py,sha256=ypdXvC6oJlsUGm4PiTJWXrtTi9w0K9IpoO0-8gM1hZ4,11295
 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,7 +41,7 @@ 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=_QoxSDMW6nmpH82Un2LDsVIBg9KIx8npRwSjY4TEGYA,31830
+xax/task/mixins/train.py,sha256=Fgx2SWGC0e1QtRv1iTXXNg45dzbCzur3UHvjRZfOoiM,33465
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
 xax/utils/experiments.py,sha256=bj8BftSHT3fFzfiJ0Co0WvqWo0rUS8kQnQYpVvH8FTM,29942
@@ -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=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.13.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.13.dist-info/METADATA,sha256=-foHRw3ph7yxBmMmjO_oqZqwvdEROYTH4Drc9P58ujI,1880
-xax-0.2.13.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
-xax-0.2.13.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.13.dist-info/RECORD,,
+xax-0.2.14.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.14.dist-info/METADATA,sha256=5Agik9rI2VgDzlElHUVYpQm8JVOytLrcIxtNMCa_UUE,1880
+xax-0.2.14.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+xax-0.2.14.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.14.dist-info/RECORD,,