xax 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

xax/task/mixins/train.py

@@ -13,14 +13,24 @@ import traceback
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, is_dataclass
 from threading import Thread
-from typing import Any, Generic, Literal, Mapping, Sequence, TypeVar, cast, get_args
+from typing import (
+    Any,
+    Generator,
+    Generic,
+    Iterator,
+    Literal,
+    Mapping,
+    Sequence,
+    TypeVar,
+    cast,
+    get_args,
+)
 
 import equinox as eqx
 import jax
-import jax.numpy as jnp
 import numpy as np
 import optax
-from jaxtyping import Array, PyTree
+from jaxtyping import Array, PRNGKeyArray, PyTree
 from omegaconf import DictConfig
 
 from xax.core.conf import field
@@ -130,6 +140,7 @@ class ValidStepTimer:
         return False
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class TrainConfig(
     CheckpointingConfig,
@@ -191,16 +202,13 @@ class TrainMixin(
         # The kind of step that was specified in the config.
         self._step_kind = cast_step_kind(self.config.step_kind)
 
-    def prng_key(self) -> jnp.ndarray:
+    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,
-            },
-        )
+        state.elapsed_time_s = time.time() - state.start_time_s
+        return state
 
     def log_train_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
         """Override this function to do logging during the training phase.
@@ -228,16 +236,19 @@ class TrainMixin(
             state: The current training state.
         """
 
-    def log_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
-        phase = state.phase
-
-        # Log the state timers.
-        timer = self.state_timers[phase]
+    def log_state_timers(self, state: State) -> None:
+        timer = self.state_timers[state.phase]
         timer.step(state)
         for ns, d in timer.log_dict().items():
             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:
+        phase = state.phase
+
+        self.logger.log_scalar("loss", loss, namespace="loss")
+        self.log_state_timers(state)
+
         # Delegate to the appropriate logging function based on the phase.
         match phase:
             case "train":
@@ -247,8 +258,10 @@ class TrainMixin(
             case _:
                 raise KeyError(f"Unknown phase: {phase}")
 
+        self.write_logs(state)
+
     @abstractmethod
-    def get_model(self) -> PyTree:
+    def get_model(self, key: PRNGKeyArray) -> PyTree:
         """Returns the Equinox model to train.
 
         Returns:
@@ -266,7 +279,10 @@ class TrainMixin(
     def get_initial_opt_state(self, model: PyTree, optimizer: optax.GradientTransformation) -> optax.OptState:
         return optimizer.init(eqx.filter(model, eqx.is_array))
 
-    def load_initial_state(self) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State]:
+    def load_initial_state(
+        self,
+        key: PRNGKeyArray,
+    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State]:
         init_ckpt_path = self.get_init_ckpt_path()
 
         if init_ckpt_path is not None:
@@ -279,18 +295,18 @@ class TrainMixin(
                 return model, optimizer, opt_state, state
 
         with self.step_context("get_model"):
-            model = self.get_model()
+            model = self.get_model(key)
 
         with self.step_context("get_optimizer"):
             optimizer = self.get_optimizer()
 
         with self.step_context("get_initial_opt_state"):
-            opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
+            opt_state = self.get_initial_opt_state(model, optimizer)
 
         return model, optimizer, opt_state, State.init_state()
 
-    @abstractmethod
-    def get_output(self, model: PyTree, batch: Batch, state: State) -> Output:
+    @eqx.filter_jit
+    def get_output(self, model: PyTree, batch: Batch) -> Output:
         """Gets the output from the model.
 
         By default, we assume the model is a function that takes the batch as
@@ -300,10 +316,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")
 
-    def compute_loss(self, model: PyTree, batch: Batch, output: Output, state: State) -> Array:
+    @eqx.filter_jit
+    def compute_loss(self, model: PyTree, batch: Batch, output: Output) -> Array:
         """Gets the loss for the current batch.
 
         By default, we assume the model is a function that takes the batch as
@@ -314,7 +331,6 @@ 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.
@@ -323,9 +339,10 @@ class TrainMixin(
             raise ValueError(f"When model output is not the loss, you must override `compute_loss`. Got {type(output)}")
         return output
 
-    def get_output_and_loss(self, model: PyTree, batch: Batch, state: State) -> tuple[Array, Output]:
-        output = self.get_output(model, batch, state)
-        loss = self.compute_loss(model, batch, output, state)
+    @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
 
     @eqx.filter_jit
@@ -335,10 +352,9 @@ class TrainMixin(
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
-        state: State,
     ) -> tuple[Array, PyTree, optax.OptState, Output]:
-        (loss, output), grads = eqx.filter_value_and_grad(self.get_output_and_loss, has_aux=True)(model, batch, state)
-        updates, opt_state = optimizer.update(grads, opt_state)
+        (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
 
@@ -377,7 +393,13 @@ class TrainMixin(
         self._last_printed_remaining_time = state.elapsed_time_s
         remaining_seconds = remaining_percent * state.elapsed_time_s / (1 - remaining_percent)
         termination_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time() + remaining_seconds))
-        logger.info("Estimated finish time: %s", termination_time)
+        # logger.info("Estimated finish time: %s", termination_time)
+        jax.debug.print("Estimated finish time: {}", termination_time)
+
+    def get_remaining_percent(self, state: State) -> float | None:
+        if self.config.max_steps is None:
+            return None
+        return (self.config.max_steps - self.get_step(state)) / self.config.max_steps
 
     def is_training_over(self, state: State) -> bool:
         if self._training_over_flag:
@@ -385,7 +407,6 @@ class TrainMixin(
         remaining_percent = self.get_remaining_percent(state)
         if remaining_percent is None:
             return False
-        self.logger.log_scalar("percent", remaining_percent, namespace="⏰ remaining")
         self.maybe_log_termination_time(remaining_percent, state)
         return remaining_percent <= 0.0
 
@@ -400,59 +421,124 @@ class TrainMixin(
             case _:
                 raise ValueError(f"Invalid step kind {self._step_kind}")
 
-    def get_remaining_percent(self, state: State) -> float | None:
-        if self.config.max_steps is None:
-            return None
-        return (self.config.max_steps - self.get_step(state)) / self.config.max_steps
-
     def log_state(self) -> None:
         logger.log(LOG_STATUS, self.task_path)
         logger.log(LOG_STATUS, self.task_name)
-        self.logger.log_git_state(get_git_state(self))
-        self.logger.log_training_code(get_training_code(self))
-        self.logger.log_config(cast(DictConfig, self.config))
+        self.logger.log_file("git_state.txt", get_git_state(self))
+        self.logger.log_file("training_code.txt", get_training_code(self))
+        self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
 
+    @eqx.filter_jit
     def train_step(
         self,
         model: 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
+
+    @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
+
+    def train_loop(
+        self,
+        model: PyTree,
+        optimizer: optax.GradientTransformation,
+        opt_state: optax.OptState,
+        train_pf: Iterator[Batch],
+        valid_pf: Iterator[Batch],
         state: State,
-    ) -> tuple[PyTree, optax.OptState, State]:
-        state = state.with_phase("train")
-        loss, model, opt_state, output = self.update(model, optimizer, opt_state, batch, state)
-        self.logger.log_scalar("loss", loss, namespace="loss")
-        self.log_step(model, batch, output, state)
-        self.write_logs(state)
-        return (
-            model,
-            opt_state,
-            state.replace(
-                {
-                    "num_steps": state.num_steps + 1,
-                    "num_samples": state.num_samples + (self.get_size_of_batch(batch) or 0),
-                },
-            ),
-        )
+    ) -> None:
+        while not self.is_training_over(state):
+            if self.valid_step_timer.is_valid_step(state):
+                valid_batch = next(valid_pf)
+                model, loss, output = self.val_step(model, valid_batch)
+
+                # 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
+
+            with self.step_context("on_step_start"):
+                state = self.on_step_start(state)
+
+            with self.step_context("update_state"):
+                train_batch = next(train_pf)
+                model, opt_state, loss, output = self.train_step(model, optimizer, opt_state, train_batch)
+
+            # Perform logging.
+            with self.step_context("write_logs"):
+                state.phase = "train"
+                self.log_step(model, train_batch, output, loss, state)
+                state.num_steps += 1
+                state.num_samples += self.get_size_of_batch(train_batch) or 0
+
+            with self.step_context("on_step_end"):
+                state = self.on_step_end(state)
+
+            if self.should_checkpoint(state):
+                self.save_checkpoint(model, optimizer, opt_state, state)
 
-    def val_step(self, model: PyTree, batch: Batch, state: State) -> tuple[PyTree, State]:
-        state = state.with_phase("valid")
-        loss, output = eqx.filter_jit(self.get_output_and_loss)(model, batch, state)
-        self.logger.log_scalar("loss", loss, namespace="loss")
-        self.log_step(model, batch, output, state)
-        self.write_logs(state)
-        return model, state.replace(
-            {
-                "num_valid_steps": state.num_valid_steps + 1,
-                "num_valid_samples": state.num_valid_samples + (self.get_size_of_batch(batch) or 0),
-            },
-        )
+        # After finishing training, save the final checkpoint.
+        self.save_checkpoint(model, optimizer, opt_state, state)
+
+    @contextlib.contextmanager
+    def get_train_iterator(self) -> Generator[Iterator[Batch], None, None]:
+        try:
+            train_iterator: Iterator[Batch] = self.get_data_iterator("train")
+            yield train_iterator
+            return
+        except NotImplementedError:
+            pass
+
+        with self.step_context("get_dataset"):
+            train_ds = self.get_dataset("train")
+
+        with self.step_context("get_dataloader"):
+            train_dl = self.get_dataloader(train_ds, "train")
+
+        with self.step_context("get_prefetcher"):
+            train_pf = self.get_prefetcher(train_dl)
+
+        try:
+            with train_pf as train_pf_ctx:
+                yield train_pf_ctx
+        finally:
+            logger.info("Closing train prefetcher")
+
+    @contextlib.contextmanager
+    def get_valid_iterator(self) -> Generator[Iterator[Batch], None, None]:
+        try:
+            valid_iterator: Iterator[Batch] = self.get_data_iterator("valid")
+            yield valid_iterator
+            return
+        except NotImplementedError:
+            pass
+
+        with self.step_context("get_dataset"):
+            valid_ds = self.get_dataset("valid")
+
+        with self.step_context("get_dataloader"):
+            valid_dl = self.get_dataloader(valid_ds, "valid")
+
+        with self.step_context("get_prefetcher"):
+            valid_pf = self.get_prefetcher(valid_dl)
+
+        try:
+            with valid_pf as valid_pf_ctx:
+                yield valid_pf_ctx
+        finally:
+            logger.info("Closing valid prefetcher")
 
     def run(self) -> None:
-        self.run_training_loop()
+        self.run_training()
 
-    def run_training_loop(self) -> None:
+    def run_training(self) -> None:
         """Runs the training loop.
 
         Args:
@@ -464,33 +550,16 @@ class TrainMixin(
         Raises:
             ValueError: If the task is not a supervised learning task
         """
-        with contextlib.ExitStack() as ctx:
+        with self:
+            key = self.prng_key()
+
             self.set_loggers()
 
             if is_master():
                 Thread(target=self.log_state, daemon=True).start()
 
-            # Gets the datasets.
-            with self.step_context("get_dataset"):
-                train_ds = self.get_dataset("train")
-                valid_ds = self.get_dataset("valid")
-
-            # Gets the dataloaders.
-            with self.step_context("get_dataloader"):
-                train_dl = self.get_dataloader(train_ds, "train")
-                valid_dl = self.get_dataloader(valid_ds, "valid")
-
-            # Gets the prefetchers.
-            with self.step_context("get_prefetcher"):
-                train_pf = self.get_prefetcher(train_dl)
-                valid_pf = self.get_prefetcher(valid_dl)
-
-            ctx.enter_context(self)
-            ctx.enter_context(train_pf)
-            ctx.enter_context(valid_pf)
-
-            model, optimizer, opt_state, state = self.load_initial_state()
-
+            key, model_key = jax.random.split(key)
+            model, optimizer, opt_state, state = self.load_initial_state(model_key)
             state = self.on_training_start(state)
 
             def on_exit() -> None:
@@ -499,43 +568,34 @@ class TrainMixin(
             # Handle user-defined interrupts during the training loop.
             self.add_signal_handler(on_exit, signal.SIGUSR1, signal.SIGTERM)
 
-            try:
-                while True:
-                    while True:
-                        if self.is_training_over(state):
-                            raise TrainingFinishedError
-
-                        if self.valid_step_timer.is_valid_step(state):
-                            model, state = self.val_step(model, next(valid_pf), state)
-
-                        with self.step_context("on_step_start"):
-                            state = self.on_step_start(state)
-
-                        model, opt_state, state = self.train_step(model, optimizer, opt_state, next(train_pf), state)
-
-                        with self.step_context("on_step_end"):
-                            state = self.on_step_end(state)
-
-                        if self.should_checkpoint(state):
-                            self.save_checkpoint(model, optimizer, opt_state, state)
-
-            except TrainingFinishedError:
-                if is_master():
-                    show_info(
-                        f"Finished training after {state.num_steps} steps, {state.num_samples} samples",
-                        important=True,
+            with self.get_train_iterator() as train_pf, self.get_valid_iterator() as valid_pf:
+                try:
+                    self.train_loop(
+                        model=model,
+                        optimizer=optimizer,
+                        opt_state=opt_state,
+                        train_pf=train_pf,
+                        valid_pf=valid_pf,
+                        state=state,
                     )
-                self.save_checkpoint(model, optimizer, opt_state, state)
-
-            except (KeyboardInterrupt, bdb.BdbQuit):
-                if is_master():
-                    show_info("Interrupted training", important=True)
-
-            except BaseException:
-                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, optimizer, opt_state, state)
 
-            finally:
-                state = self.on_training_end(state)
+                except TrainingFinishedError:
+                    if is_master():
+                        show_info(
+                            f"Finished training after {state.num_steps} steps, {state.num_samples} samples",
+                            important=True,
+                        )
+                    self.save_checkpoint(model, optimizer, opt_state, state)
+
+                except (KeyboardInterrupt, bdb.BdbQuit):
+                    if is_master():
+                        show_info("Interrupted training", important=True)
+
+                except BaseException:
+                    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, optimizer, opt_state, state)
+
+                finally:
+                    state = self.on_training_end(state)

xax/task/script.py

@@ -22,7 +22,7 @@ from xax.task.mixins import (
 )
 
 
-@dataclass
+@dataclass(kw_only=True)
 class ScriptConfig(
     CPUStatsConfig,
     GPUStatsConfig,

xax/task/task.py

@@ -3,12 +3,16 @@
 from dataclasses import dataclass
 from typing import Generic, TypeVar
 
+import jax
+
 from xax.task.base import BaseConfig, BaseTask
 from xax.task.mixins import (
     ArtifactsConfig,
     ArtifactsMixin,
     CheckpointingConfig,
     CheckpointingMixin,
+    CompileConfig,
+    CompileMixin,
     CPUStatsConfig,
     CPUStatsMixin,
     DataloadersConfig,
@@ -28,10 +32,12 @@ from xax.task.mixins import (
 )
 
 
+@jax.tree_util.register_dataclass
 @dataclass
 class Config(
     TrainConfig,
     CheckpointingConfig,
+    CompileConfig,
     DataloadersConfig,
     CPUStatsConfig,
     GPUStatsConfig,
@@ -51,6 +57,7 @@ ConfigT = TypeVar("ConfigT", bound=Config)
 class Task(
     TrainMixin[ConfigT],
     CheckpointingMixin[ConfigT],
+    CompileMixin[ConfigT],
     DataloadersMixin[ConfigT],
     CPUStatsMixin[ConfigT],
     GPUStatsMixin[ConfigT],

xax/utils/tensorboard.py

@@ -2,10 +2,14 @@
 
 import functools
 import io
+import os
+import tempfile
 import time
 from pathlib import Path
 from typing import Literal, TypedDict
 
+import numpy as np
+import PIL.Image
 from PIL.Image import Image as PILImage
 from tensorboard.compat.proto.config_pb2 import RunMetadata
 from tensorboard.compat.proto.event_pb2 import Event, TaggedRunMetadata
@@ -186,6 +190,50 @@ class TensorboardWriter:
             walltime=walltime,
         )
 
+    def add_video(
+        self,
+        tag: str,
+        value: np.ndarray,
+        global_step: int | None = None,
+        walltime: float | None = None,
+        fps: int = 30,
+    ) -> None:
+        assert value.ndim == 4, "Video must be 4D array (T, H, W, C)"
+        images = [PIL.Image.fromarray(frame) for frame in value]
+
+        # Create temporary file for GIF
+        temp_file = tempfile.NamedTemporaryFile(suffix=".gif", delete=False)
+        try:
+            images[0].save(temp_file.name, save_all=True, append_images=images[1:], duration=int(1000 / fps), loop=0)
+            with open(temp_file.name, "rb") as f:
+                video_string = f.read()
+
+        finally:
+            # Clean up temporary file
+            try:
+                os.remove(temp_file.name)
+            except OSError:
+                pass
+
+        # Add to summary
+        self.pb_writer.add_summary(
+            Summary(
+                value=[
+                    Summary.Value(
+                        tag=tag,
+                        image=Summary.Image(
+                            height=value.shape[1],
+                            width=value.shape[2],
+                            colorspace=value.shape[3],
+                            encoded_image_string=video_string,
+                        ),
+                    ),
+                ],
+            ),
+            global_step=global_step,
+            walltime=walltime,
+        )
+
     def add_text(
         self,
         tag: str,

{xax-0.0.5.dist-info → xax-0.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: xax
-Version: 0.0.5
+Version: 0.0.6
 Summary: The xax project
 Home-page: https://github.com/dpshai/xax
 Author: Benjamin Bolte
@@ -12,6 +12,8 @@ Requires-Dist: jaxtyping
 Requires-Dist: equinox
 Requires-Dist: optax
 Requires-Dist: dpshdl
+Requires-Dist: chex
+Requires-Dist: importlib-resources
 Requires-Dist: cloudpickle
 Requires-Dist: pillow
 Requires-Dist: omegaconf
@@ -28,6 +30,14 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
+Dynamic: author
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 # xax
 

xax-0.0.6.dist-info/RECORD

@@ -0,0 +1,52 @@
+xax/__init__.py,sha256=RTUsDh_R0TFa09q-_U0vd-eCYRC-bCaHqHlayp8U2hU,9736
+xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
+xax/requirements.txt,sha256=NmU9PNJhfLtNqqtWWf8WqMjgbBPCn_yt8oMGAgS7Fno,291
+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/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
+xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
+xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
+xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/task/base.py,sha256=LHDmM2c_Ps5cGEzn_QUpmyInD7zJJm3Yt9eSeij2Vus,7297
+xax/task/logger.py,sha256=orN1jmM4SIR2EiYk8bNoJZscmhX1FytADBU6p9qpows,29256
+xax/task/script.py,sha256=4LyXrpj0V36TjAZT4lvQeiOTqa5U2tommHKwgWDCE24,1025
+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
+xax/task/launchers/cli.py,sha256=cK7Nm-3fO-W2gTxpn3FEThsT2NvneS2w0UjA1Nt-84A,1402
+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=yXHb1bmfsEnk4p0F1Up1ertWYdcPAFZm25NT8wE3Jb8,4045
+xax/task/loggers/state.py,sha256=6bG-NRsSUzAukYiglCT0oDj8zRMpffH4e1TKWGw1x4k,959
+xax/task/loggers/stdout.py,sha256=nxQXkS9JUR38RKsU9qj0dgePKguK0BFa9nl_BdGO8cE,6758
+xax/task/loggers/tensorboard.py,sha256=FGW96z77oG0Kf3cO6Zznx5U3kJNzPWcuSkpY4RnbFCo,6909
+xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
+xax/task/mixins/artifacts.py,sha256=1H7ZbR-KSsXhVtqGVlqMi-TXfn1-dM7YnTCLVuw594s,3835
+xax/task/mixins/checkpointing.py,sha256=AMlobojybvJdDZcNCxm1DHSVC_2Qvnu_MbRcsc_8eoA,8508
+xax/task/mixins/compile.py,sha256=FRsxwLnZjjxpeWJ7Bx_d8XUY50oDoGidgpeRt4ejeQk,3377
+xax/task/mixins/cpu_stats.py,sha256=C_t71UTrv4LwQzhO5iubsfomj4jYa9bzpE4zBcHdoHM,9211
+xax/task/mixins/data_loader.py,sha256=WjMWk9uACfBMMClLMcLPkE0WNIvlCZnmqyyqLqJpjX0,6545
+xax/task/mixins/gpu_stats.py,sha256=IGPBro9xzSivwD43zM18lWcuei7IhA8LilxSPHqNl4I,8747
+xax/task/mixins/logger.py,sha256=CIQ4w4K3FcxN6A9xUfITdVkulSxPa4iaTe6cbs9ruaM,1958
+xax/task/mixins/process.py,sha256=d1opVgvc6bOFXb7R58b07F4P5lbSZIzYaajtE0eBbpw,1477
+xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
+xax/task/mixins/step_wrapper.py,sha256=DJw42mUGwgKx2tkeqatKR9_F4J8ug4wmxKMeJPmhcVQ,1560
+xax/task/mixins/train.py,sha256=dhGL_IuDaJy39BooYlO7JO-_EotKldtBhBplDGU_AnM,21745
+xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/utils/experiments.py,sha256=qT3H0fyVH8DN417x7T0Xmz4SKoogW81-EHcZfyktFI8,28300
+xax/utils/jax.py,sha256=VzEVB766UyH3_cgN6UP0FkCsDuGlYg5KJj8YJS4yYUk,439
+xax/utils/logging.py,sha256=ST1hp2C2xntVVJBUHwo3YxPK19fBLNvHU2WGO1xqcXA,6418
+xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
+xax/utils/tensorboard.py,sha256=oGq2E3Yr0z2xaACv2UOVt_CHEVc8fBxI8V1M99Fd34E,9742
+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-0.0.6.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.0.6.dist-info/METADATA,sha256=YO2c2PUMWkH1ILfPhFWKK4Sodbo9qUpUOCIkm4aLHfg,1171
+xax-0.0.6.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+xax-0.0.6.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.0.6.dist-info/RECORD,,

{xax-0.0.5.dist-info → xax-0.0.6.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.6.0)
+Generator: setuptools (75.8.2)
 Root-Is-Purelib: true
 Tag: py3-none-any