xax 0.1.16__py3-none-any.whl → 0.2.1__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.16"
+__version__ = "0.2.1"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -44,6 +44,7 @@ __all__ = [
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "quat_to_euler",
+    "quat_to_rotmat",
     "rotate_vector_by_quat",
     "cross_entropy",
     "cast_norm_type",
@@ -206,6 +207,7 @@ NAME_MAP: dict[str, str] = {
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "quat_to_euler": "nn.geom",
+    "quat_to_rotmat": "nn.geom",
     "rotate_vector_by_quat": "nn.geom",
     "cross_entropy": "nn.losses",
     "cast_norm_type": "nn.norm",
@@ -369,6 +371,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
         quat_to_euler,
+        quat_to_rotmat,
         rotate_vector_by_quat,
     )
     from xax.nn.losses import cross_entropy

xax/core/state.py

@@ -12,6 +12,14 @@ from xax.core.conf import field
 Phase = Literal["train", "valid"]
 
 
+def _phase_to_int(phase: Phase) -> int:
+    return {"train": 0, "valid": 1}[phase]
+
+
+def _int_to_phase(i: int) -> Phase:
+    return cast(Phase, ["train", "valid"][i])
+
+
 class StateDict(TypedDict, total=False):
     num_steps: NotRequired[int]
     num_samples: NotRequired[int]
@@ -35,7 +43,7 @@ class State:
 
     @property
     def phase(self) -> Phase:
-        return cast(Phase, ["train", "valid"][self._phase])
+        return _int_to_phase(self._phase)
 
     @classmethod
     def init_state(cls) -> "State":
@@ -74,3 +82,20 @@ class State:
                 case _:
                     raise ValueError(f"Invalid phase: {phase}")
         return State(**{**asdict(self), **kwargs, **extra_kwargs})
+
+    def to_dict(self) -> dict[str, int | float | str]:
+        return {
+            "num_steps": int(self.num_steps),
+            "num_samples": int(self.num_samples),
+            "num_valid_steps": int(self.num_valid_steps),
+            "num_valid_samples": int(self.num_valid_samples),
+            "start_time_s": float(self.start_time_s),
+            "elapsed_time_s": float(self.elapsed_time_s),
+            "phase": str(self.phase),
+        }
+
+    @classmethod
+    def from_dict(cls, d: dict[str, int | float | str]) -> "State":
+        if "phase" in d:
+            d["_phase"] = _phase_to_int(cast(Phase, d.pop("phase")))
+        return cls(**d)  # type: ignore[arg-type]

xax/nn/geom.py

@@ -177,3 +177,37 @@ def cubic_bezier_interpolation(y_start: Array, y_end: Array, x: Array) -> Array:
     y_diff = y_end - y_start
     bezier = x**3 + 3 * (x**2 * (1 - x))
     return y_start + y_diff * bezier
+
+
+def quat_to_rotmat(quat: Array, eps: float = 1e-6) -> Array:
+    """Converts a quaternion to a rotation matrix.
+
+    Args:
+        quat: The quaternion to convert, shape (*, 4).
+        eps: A small epsilon value to avoid division by zero.
+
+    Returns:
+        The rotation matrix, shape (*, 3, 3).
+    """
+    quat = quat / (jnp.linalg.norm(quat, axis=-1, keepdims=True) + eps)
+    w, x, y, z = jnp.split(quat, 4, axis=-1)
+
+    xx = 1 - 2 * (y * y + z * z)
+    xy = 2 * (x * y - z * w)
+    xz = 2 * (x * z + y * w)
+    yx = 2 * (x * y + z * w)
+    yy = 1 - 2 * (x * x + z * z)
+    yz = 2 * (y * z - x * w)
+    zx = 2 * (x * z - y * w)
+    zy = 2 * (y * z + x * w)
+    zz = 1 - 2 * (x * x + y * y)
+
+    # Corrected stacking: row-major order
+    return jnp.concatenate(
+        [
+            jnp.concatenate([xx, xy, xz], axis=-1)[..., None, :],
+            jnp.concatenate([yx, yy, yz], axis=-1)[..., None, :],
+            jnp.concatenate([zx, zy, zz], axis=-1)[..., None, :],
+        ],
+        axis=-2,
+    )

xax/requirements.txt

@@ -2,16 +2,16 @@
 
 # Core ML/JAX dependencies
 attrs
+chex
+dpshdl
+equinox
+importlib-resources
 jax
 jaxtyping
-equinox
 optax
-dpshdl
-chex
-importlib-resources
+orbax-checkpoint
 
 # Data processing and serialization
-cloudpickle
 pillow
 
 # Configuration and project management

xax/task/base.py

@@ -79,7 +79,7 @@ class BaseTask(Generic[Config]):
     def on_training_end(self, state: State) -> State:
         return state
 
-    def on_after_checkpoint_save(self, ckpt_path: Path, state: State) -> State:
+    def on_after_checkpoint_save(self, ckpt_path: Path, state: State | None) -> State | None:
         return state
 
     @functools.cached_property

xax/task/logger.py

@@ -18,11 +18,22 @@ from abc import ABC, abstractmethod
 from collections import defaultdict
 from dataclasses import dataclass
 from types import TracebackType
-from typing import Callable, Iterator, Literal, Self, Sequence, TypeVar, cast, get_args
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    Literal,
+    Self,
+    Sequence,
+    TypeVar,
+    cast,
+    get_args,
+)
 
 import jax
 import jax.numpy as jnp
 import numpy as np
+from jax._src.core import ClosedJaxpr
 from jaxtyping import Array
 from PIL import Image, ImageDraw, ImageFont
 from PIL.Image import Image as PILImage
@@ -194,7 +205,10 @@ def tile_images(images: list[PILImage], sep: int = 0) -> PILImage:
     return tiled
 
 
-def as_numpy(array: Array) -> np.ndarray:
+def as_numpy(array: Array | np.ndarray) -> np.ndarray:
+    """Convert a JAX array or numpy array to numpy array."""
+    if isinstance(array, np.ndarray):
+        return array
     array = jax.device_get(array)
     if jax.dtypes.issubdtype(array.dtype, jnp.floating):
         array = array.astype(jnp.float32)
@@ -205,6 +219,13 @@ def as_numpy(array: Array) -> np.ndarray:
     return np.array(array)
 
 
+def as_numpy_opt(array: Array | np.ndarray | None) -> np.ndarray | None:
+    """Convert an optional JAX array or numpy array to numpy array."""
+    if array is None:
+        return None
+    return as_numpy(array)
+
+
 @dataclass(kw_only=True)
 class LogString:
     value: str
@@ -252,6 +273,19 @@ class LogHistogram:
     bucket_counts: list[int]
 
 
+@dataclass(kw_only=True)
+class LogMesh:
+    vertices: np.ndarray
+    colors: np.ndarray | None
+    faces: np.ndarray | None
+    config_dict: dict[str, Any] | None  # noqa: ANN401
+
+
+@dataclass(kw_only=True)
+class LogGraph:
+    computation: ClosedJaxpr
+
+
 @dataclass(kw_only=True)
 class LogLine:
     state: State
@@ -261,6 +295,7 @@ class LogLine:
     strings: dict[str, dict[str, LogString]]
     images: dict[str, dict[str, LogImage]]
     videos: dict[str, dict[str, LogVideo]]
+    meshes: dict[str, dict[str, LogMesh]]
 
 
 @dataclass(kw_only=True)
@@ -533,6 +568,7 @@ class Logger:
         self.strings: dict[str, dict[str, Callable[[], LogString]]] = defaultdict(dict)
         self.images: dict[str, dict[str, Callable[[], LogImage]]] = defaultdict(dict)
         self.videos: dict[str, dict[str, Callable[[], LogVideo]]] = defaultdict(dict)
+        self.meshes: dict[str, dict[str, Callable[[], LogMesh]]] = defaultdict(dict)
         self.default_namespace = default_namespace
         self.loggers: list[LoggerImpl] = []
 
@@ -560,6 +596,7 @@ class Logger:
             strings={k: {kk: v() for kk, v in v.items()} for k, v in self.strings.items()},
             images={k: {kk: v() for kk, v in v.items()} for k, v in self.images.items()},
             videos={k: {kk: v() for kk, v in v.items()} for k, v in self.videos.items()},
+            meshes={k: {kk: v() for kk, v in v.items()} for k, v in self.meshes.items()},
         )
 
     def clear(self) -> None:
@@ -569,6 +606,7 @@ class Logger:
         self.strings.clear()
         self.images.clear()
         self.videos.clear()
+        self.meshes.clear()
 
     def write(self, state: State) -> None:
         """Writes the current step's logging information.
@@ -1051,6 +1089,73 @@ class Logger:
 
         self.videos[namespace][key] = video_future
 
+    def log_mesh(
+        self,
+        key: str,
+        vertices: np.ndarray | Array | Callable[[], np.ndarray | Array],
+        colors: np.ndarray | Array | None | Callable[[], np.ndarray | Array | None] = None,
+        faces: np.ndarray | Array | None | Callable[[], np.ndarray | Array | None] = None,
+        config_dict: dict[str, Any] | None = None,
+        *,
+        namespace: str | None = None,
+    ) -> None:
+        if not self.active:
+            raise RuntimeError("The logger is not active")
+        namespace = self.resolve_namespace(namespace)
+
+        @functools.lru_cache(maxsize=None)
+        def mesh_future() -> LogMesh:
+            with ContextTimer() as timer:
+                # Get the raw values
+                vertices_val = vertices() if callable(vertices) else vertices
+                colors_val = colors() if callable(colors) else colors
+                faces_val = faces() if callable(faces) else faces
+
+                # Convert to numpy arrays with proper type handling
+                vertices_np = as_numpy(vertices_val)
+                colors_np = as_numpy_opt(colors_val)
+                faces_np = as_numpy_opt(faces_val)
+
+                # Checks vertices shape.
+                if vertices_np.ndim == 2:
+                    vertices_np = vertices_np[None]
+                if vertices_np.shape[-1] != 3 or vertices_np.ndim != 3:
+                    raise ValueError("Vertices must have shape (N, 3) or (B, N, 3)")
+
+                # Checks colors shape.
+                if colors_np is not None:
+                    if colors_np.ndim == 2:
+                        colors_np = colors_np[None]
+                    if colors_np.shape[-1] != 3 or colors_np.ndim != 3:
+                        raise ValueError("Colors must have shape (N, 3) or (B, N, 3)")
+
+                # Checks faces shape.
+                if faces_np is not None:
+                    if faces_np.ndim == 2:
+                        faces_np = faces_np[None]
+                    if faces_np.shape[-1] != 3 or faces_np.ndim != 3:
+                        raise ValueError("Faces must have shape (N, 3) or (B, N, 3)")
+
+                # Ensures colors dtype is uint8.
+                if colors_np is not None:
+                    if colors_np.dtype != np.uint8:
+                        colors_np = (colors_np * 255).astype(np.uint8)
+
+                # Ensures faces dtype is int32.
+                if faces_np is not None:
+                    if faces_np.dtype != np.int32:
+                        faces_np = faces_np.astype(np.int32)
+
+            logger.debug("Mesh Key: %s, Time: %s", key, timer.elapsed_time)
+            return LogMesh(
+                vertices=vertices_np,
+                colors=colors_np,
+                faces=faces_np,
+                config_dict=config_dict,
+            )
+
+        self.meshes[namespace][key] = mesh_future
+
     def __enter__(self) -> Self:
         self.active = True
         for logger in self.loggers:

xax/task/loggers/tensorboard.py

@@ -70,6 +70,9 @@ class TensorboardLogger(LoggerImpl):
         self._started = True
 
     def worker_thread(self) -> None:
+        if os.environ.get("DISABLE_TENSORBOARD", "0") == "1":
+            return
+
         time.sleep(self.wait_seconds)
 
         port = int(os.environ.get("TENSORBOARD_PORT", DEFAULT_TENSORBOARD_PORT))
@@ -213,6 +216,19 @@ class TensorboardLogger(LoggerImpl):
                     video_value.frames,
                     fps=video_value.fps,
                     global_step=line.state.num_steps,
+                    walltime=walltime,
+                )
+
+        for namespace, meshes in line.meshes.items():
+            for mesh_key, mesh_value in meshes.items():
+                writer.add_mesh(
+                    f"{namespace}/{mesh_key}",
+                    vertices=mesh_value.vertices,
+                    faces=mesh_value.faces,
+                    colors=mesh_value.colors,
+                    config_dict=mesh_value.config_dict,
+                    global_step=line.state.num_steps,
+                    walltime=walltime,
                 )
 
         for name, contents in self.files.items():

xax/task/mixins/checkpointing.py

@@ -6,9 +6,9 @@ import logging
 import tarfile
 from dataclasses import asdict, dataclass
 from pathlib import Path
-from typing import Any, Callable, Generic, Literal, TypeVar, cast, overload
+from typing import Generic, Literal, TypeVar, cast, overload
 
-import cloudpickle
+import equinox as eqx
 import jax
 import optax
 from jaxtyping import PyTree
@@ -63,8 +63,7 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
 
     def get_init_ckpt_path(self) -> Path | None:
         if self._exp_dir is not None:
-            ckpt_path = self.get_ckpt_path()
-            if ckpt_path.exists():
+            if (ckpt_path := self.get_ckpt_path()).exists():
                 return ckpt_path
         if self.config.load_from_ckpt_path is not None:
             ckpt_path = Path(self.config.load_from_ckpt_path)
@@ -84,93 +83,129 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
         return False
 
     @overload
-    def load_checkpoint(
+    def load_ckpt_with_template(
         self,
         path: Path,
-        part: Literal["all"] = "all",
-    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State, DictConfig]: ...
+        *,
+        part: Literal["all"],
+        model_template: PyTree,
+        optimizer_template: PyTree,
+        opt_state_template: PyTree,
+    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State, Config]: ...
 
     @overload
-    def load_checkpoint(
+    def load_ckpt_with_template(
         self,
         path: Path,
-        part: Literal["model_state_config"] = "model_state_config",
-    ) -> tuple[PyTree, State, DictConfig]: ...
+        *,
+        part: Literal["model_state_config"],
+        model_template: PyTree,
+    ) -> tuple[PyTree, State, Config]: ...
 
     @overload
-    def load_checkpoint(
+    def load_ckpt_with_template(
         self,
         path: Path,
+        *,
         part: Literal["model"],
+        model_template: PyTree,
     ) -> PyTree: ...
 
     @overload
-    def load_checkpoint(
+    def load_ckpt_with_template(
         self,
         path: Path,
+        *,
         part: Literal["opt"],
+        optimizer_template: PyTree,
     ) -> optax.GradientTransformation: ...
 
     @overload
-    def load_checkpoint(
+    def load_ckpt_with_template(
         self,
         path: Path,
+        *,
         part: Literal["opt_state"],
+        opt_state_template: PyTree,
     ) -> optax.OptState: ...
 
     @overload
-    def load_checkpoint(
+    def load_ckpt_with_template(
         self,
         path: Path,
+        *,
         part: Literal["state"],
     ) -> State: ...
 
     @overload
-    def load_checkpoint(
+    def load_ckpt_with_template(
         self,
         path: Path,
+        *,
         part: Literal["config"],
-    ) -> DictConfig: ...
+    ) -> Config: ...
 
-    def load_checkpoint(
+    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, DictConfig]
-        | tuple[PyTree, State, DictConfig]
+        tuple[PyTree, optax.GradientTransformation, optax.OptState, State, Config]
+        | tuple[PyTree, State, Config]
         | PyTree
         | optax.GradientTransformation
         | optax.OptState
         | State
-        | DictConfig
+        | 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 cloudpickle.load(model)
+                return eqx.tree_deserialise_leaves(io.BytesIO(model.read()), model_template)
 
             def get_opt() -> optax.GradientTransformation:
-                if (opt := tar.extractfile("opt")) is None:
-                    raise ValueError(f"Checkpoint does not contain an opt file: {path}")
-                return cloudpickle.load(opt)
+                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 opt_state file: {path}")
-                return cloudpickle.load(opt_state)
+                    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(**json.loads(state.read().decode()))
 
-            def get_config() -> DictConfig:
+            def get_config() -> Config:
                 if (config := tar.extractfile("config")) is None:
                     raise ValueError(f"Checkpoint does not contain a config file: {path}")
-                return cast(DictConfig, OmegaConf.load(config))
+                return self.get_config(cast(DictConfig, OmegaConf.load(config)), use_cli=False)
 
             match part:
                 case "model":
@@ -192,51 +227,90 @@ class CheckpointingMixin(ArtifactsMixin[Config], Generic[Config]):
 
     def save_checkpoint(
         self,
-        model: PyTree,
-        optimizer: optax.GradientTransformation,
-        opt_state: optax.OptState,
-        state: State,
+        model: PyTree | None = None,
+        optimizer: optax.GradientTransformation | None = None,
+        opt_state: 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
+            aux_data: Additional data to save
+            opt_state: The optimizer state to save
+
+        Returns:
+            Path to the saved checkpoint
+        """
         ckpt_path = self.get_ckpt_path(state)
 
         if not is_master():
             return ckpt_path
 
-        # Gets the path to the last checkpoint.
+        # Gets the path to the last checkpoint
         logger.info("Saving checkpoint to %s", ckpt_path)
         last_ckpt_path = self.get_ckpt_path()
         ckpt_path.parent.mkdir(exist_ok=True, parents=True)
 
-        # Potentially removes the last checkpoint.
+        # Potentially removes the last checkpoint
         if last_ckpt_path.exists() and self.config.only_save_most_recent:
             if (base_ckpt := last_ckpt_path.resolve()).is_file():
                 base_ckpt.unlink()
 
-        # Combines all temporary files into a single checkpoint TAR file.
+        # Save the checkpoint components
         with tarfile.open(ckpt_path, "w:gz") as tar:
 
-            def add_file(name: str, write_fn: Callable[[io.BytesIO], Any]) -> None:
+            def add_file(name: str, buf: io.BytesIO) -> None:
+                tarinfo = tarfile.TarInfo(name)
+                tarinfo.size = buf.tell()
+                buf.seek(0)
+                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)
+
+            # Save optimizer using Equinox
+            if optimizer is not None:
+                with io.BytesIO() as buf:
+                    eqx.tree_serialise_leaves(buf, optimizer)
+                    add_file("optimizer", buf)
+
+            # Save optimizer state using Equinox
+            if opt_state is not None:
                 with io.BytesIO() as buf:
-                    write_fn(buf)
-                    tarinfo = tarfile.TarInfo(name)
-                    tarinfo.size = buf.tell()
-                    buf.seek(0)
-                    tar.addfile(tarinfo, buf)
-
-            add_file("model", lambda buf: cloudpickle.dump(model, buf))
-            add_file("opt", lambda buf: cloudpickle.dump(optimizer, buf))
-            add_file("opt_state", lambda buf: cloudpickle.dump(opt_state, buf))
-            add_file("state", lambda buf: buf.write(json.dumps(asdict(state), indent=2).encode()))
-            add_file("config", lambda buf: buf.write(OmegaConf.to_yaml(self.config).encode()))
-
-        # Updates the symlink to the new checkpoint.
+                    eqx.tree_serialise_leaves(buf, opt_state)
+                    add_file("opt_state", buf)
+
+            # Save aux data using Equinox.
+            if aux_data is not None:
+                with io.BytesIO() as buf:
+                    eqx.tree_serialise_leaves(buf, aux_data)
+                    add_file("aux_data", buf)
+
+            # Save state and config as JSON
+            def add_file_bytes(name: str, data: bytes) -> None:  # noqa: ANN401
+                info = tarfile.TarInfo(name=name)
+                info.size = len(data)
+                tar.addfile(info, io.BytesIO(data))
+
+            if state is not None:
+                add_file_bytes("state", json.dumps(asdict(state), indent=2).encode())
+            add_file_bytes("config", OmegaConf.to_yaml(self.config).encode())
+
+        # Updates the symlink to the new checkpoint
         last_ckpt_path.unlink(missing_ok=True)
         try:
             last_ckpt_path.symlink_to(ckpt_path.relative_to(last_ckpt_path.parent))
         except FileExistsError:
             logger.exception("Exception while trying to update %s", ckpt_path)
 
-        # Calls the base callback.
+        # Calls the base callback
         self.on_after_checkpoint_save(ckpt_path, state)
 
         return ckpt_path

xax/task/mixins/data_loader.py

@@ -9,6 +9,7 @@ import jax
 from dpshdl.dataloader import CollatedDataloaderItem, Dataloader
 from dpshdl.dataset import Dataset, ErrorHandlingDataset
 from dpshdl.prefetcher import Prefetcher
+from jaxtyping import PRNGKeyArray
 from omegaconf import II, MISSING
 
 from xax.core.conf import field, is_missing
@@ -103,7 +104,7 @@ class DataloadersMixin(ProcessMixin[Config], BaseTask[Config], Generic[Config],
             "or `get_data_iterator` to return an iterator for the given dataset."
         )
 
-    def get_data_iterator(self, phase: Phase) -> Iterator:
+    def get_data_iterator(self, phase: Phase, key: PRNGKeyArray) -> Iterator:
         raise NotImplementedError(
             "You must implement either the `get_dataset` method to return the dataset for the given phase, "
             "or `get_data_iterator` to return an iterator for the given dataset."

xax/task/mixins/train.py

@@ -11,7 +11,8 @@ import textwrap
 import time
 import traceback
 from abc import ABC, abstractmethod
-from dataclasses import dataclass, is_dataclass
+from dataclasses import asdict, dataclass, is_dataclass
+from pathlib import Path
 from threading import Thread
 from typing import (
     Any,
@@ -33,14 +34,13 @@ import jax.numpy as jnp
 import numpy as np
 import optax
 from jaxtyping import Array, PRNGKeyArray, PyTree
-from omegaconf import DictConfig
 
 from xax.core.conf import field
 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
+from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin, CheckpointPart
 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
@@ -50,11 +50,12 @@ from xax.utils.experiments import (
     TrainingFinishedError,
     diff_configs,
     get_diff_string,
+    get_info_json,
     get_state_file_string,
     get_training_code,
 )
 from xax.utils.jax import jit as xax_jit
-from xax.utils.logging import LOG_STATUS
+from xax.utils.logging import LOG_PING, LOG_STATUS
 from xax.utils.text import highlight_exception_message, show_info
 from xax.utils.types.frozen_dict import FrozenDict
 
@@ -340,20 +341,19 @@ class TrainMixin(
 
         if init_ckpt_path is not None:
             logger.info("Loading checkpoint from %s", init_ckpt_path)
-            if load_optimizer:
-                model, optimizer, opt_state, state, config = self.load_checkpoint(init_ckpt_path)
-                config_diff = get_diff_string(diff_configs(config, cast(DictConfig, self.config)))
-                if config_diff:
-                    logger.warning("Loaded config differs from current config:\n%s", config_diff)
-                return model, optimizer, opt_state, state
+            model, state, config = self.load_ckpt(init_ckpt_path, part="model_state_config")
+            config_diff = get_diff_string(diff_configs(asdict(config), asdict(self.config)))
+            if config_diff:
+                logger.warning("Loaded config differs from current config:\n%s", config_diff)
 
-            else:
-                model, state, config = self.load_checkpoint(init_ckpt_path, "model_state_config")
-                config_diff = get_diff_string(diff_configs(config, cast(DictConfig, self.config)))
-                if config_diff:
-                    logger.warning("Loaded config differs from current config:\n%s", config_diff)
+            if not load_optimizer:
                 return model, state
 
+            optimizer = self.load_ckpt(init_ckpt_path, part="opt")
+            opt_state = self.load_ckpt(init_ckpt_path, part="opt_state", model=model, optimizer=optimizer)
+            return model, optimizer, opt_state, state
+
+        logger.info("Starting a new training run")
         model = self.get_model(key)
         state = State.init_state()
 
@@ -365,6 +365,131 @@ class TrainMixin(
 
         return model, optimizer, opt_state, state
 
+    @overload
+    def load_ckpt(
+        self,
+        path: Path,
+        *,
+        part: Literal["all"],
+    ) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State, Config]: ...
+
+    @overload
+    def load_ckpt(
+        self,
+        path: Path,
+        *,
+        part: Literal["model_state_config"],
+    ) -> tuple[PyTree, State, Config]: ...
+
+    @overload
+    def load_ckpt(
+        self,
+        path: Path,
+        *,
+        part: Literal["model"],
+    ) -> PyTree: ...
+
+    @overload
+    def load_ckpt(
+        self,
+        path: Path,
+        *,
+        part: Literal["opt"],
+    ) -> optax.GradientTransformation: ...
+
+    @overload
+    def load_ckpt(
+        self,
+        path: Path,
+        *,
+        part: Literal["opt_state"],
+        model: PyTree | None = None,
+        optimizer: optax.GradientTransformation | None = None,
+    ) -> optax.OptState: ...
+
+    @overload
+    def load_ckpt(
+        self,
+        path: Path,
+        *,
+        part: Literal["state"],
+    ) -> State: ...
+
+    @overload
+    def load_ckpt(
+        self,
+        path: Path,
+        *,
+        part: Literal["config"],
+    ) -> Config: ...
+
+    def load_ckpt(
+        self,
+        path: str | Path,
+        *,
+        part: CheckpointPart = "all",
+        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
+        | State
+        | Config
+    ):
+        path = Path(path)
+
+        # This key isn't used for anything, it's just a required argument.
+        key = jax.random.PRNGKey(0)
+
+        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)
+
+            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")
+
+            case "opt":
+                optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
+                return self.load_ckpt_with_template(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)
+                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)
+                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)
+
+            case "state":
+                return self.load_ckpt_with_template(path, part="state")
+
+            case "config":
+                return self.load_ckpt_with_template(path, part="config")
+
+            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)
+                optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
+                optimizer = self.load_ckpt_with_template(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")
+                return model, optimizer, opt_state, state, config
+
+            case _:
+                raise ValueError(f"Unknown checkpoint part: {part}")
+
     def get_output(self, model: PyTree, batch: Batch, state: State) -> Output:
         """Gets the output from the model.
 
@@ -519,8 +644,7 @@ 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)
-        jax.debug.print("Estimated finish time: {}", termination_time)
+        logger.log(LOG_PING, "Estimated finish time: %s", termination_time)
 
     def get_remaining_percent(self, state: State) -> float | None:
         if self.config.max_steps is None:
@@ -554,6 +678,7 @@ class TrainMixin(
         self.logger.log_file("state.txt", get_state_file_string(self))
         self.logger.log_file("training_code.py", get_training_code(self))
         self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
+        self.logger.log_file("info.json", get_info_json())
 
     def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401
         return eqx.is_inexact_array(item)
@@ -627,16 +752,16 @@ class TrainMixin(
 
             if self.should_checkpoint(state):
                 model = eqx.combine(model_arr, model_static)
-                self.save_checkpoint(model, optimizer, opt_state, state)
+                self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
 
         # After finishing training, save the final checkpoint.
         model = eqx.combine(model_arr, model_static)
-        self.save_checkpoint(model, optimizer, opt_state, state)
+        self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
 
     @contextlib.contextmanager
-    def get_train_iterator(self) -> Generator[Iterator[Batch], None, None]:
+    def get_train_iterator(self, key: PRNGKeyArray) -> Generator[Iterator[Batch], None, None]:
         try:
-            train_iterator: Iterator[Batch] = self.get_data_iterator("train")
+            train_iterator: Iterator[Batch] = self.get_data_iterator("train", key=key)
             yield train_iterator
             return
         except NotImplementedError:
@@ -653,9 +778,9 @@ class TrainMixin(
             logger.info("Closing train prefetcher")
 
     @contextlib.contextmanager
-    def get_valid_iterator(self) -> Generator[Iterator[Batch], None, None]:
+    def get_valid_iterator(self, key: PRNGKeyArray) -> Generator[Iterator[Batch], None, None]:
         try:
-            valid_iterator: Iterator[Batch] = self.get_data_iterator("valid")
+            valid_iterator: Iterator[Batch] = self.get_data_iterator("valid", key=key)
             yield valid_iterator
             return
         except NotImplementedError:
@@ -699,12 +824,13 @@ class TrainMixin(
             state = self.on_training_start(state)
 
             def on_exit() -> None:
-                self.save_checkpoint(model, optimizer, opt_state, state)
+                self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
 
             # Handle user-defined interrupts during the training loop.
             self.add_signal_handler(on_exit, signal.SIGUSR1, signal.SIGTERM)
 
-            with self.get_train_iterator() as train_pf, self.get_valid_iterator() as valid_pf:
+            key, tkey, vkey = jax.random.split(key, 3)
+            with self.get_train_iterator(tkey) as train_pf, self.get_valid_iterator(vkey) as valid_pf:
                 try:
                     self.train_loop(
                         model=model,
@@ -721,7 +847,7 @@ class TrainMixin(
                             f"Finished training after {state.num_steps} steps, {state.num_samples} samples",
                             important=True,
                         )
-                    self.save_checkpoint(model, optimizer, opt_state, state)
+                    self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
 
                 except (KeyboardInterrupt, bdb.BdbQuit):
                     if is_master():
@@ -731,7 +857,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, optimizer, opt_state, state)
+                    self.save_checkpoint(model=model, optimizer=optimizer, opt_state=opt_state, state=state)
 
                 finally:
                     state = self.on_training_end(state)

xax/utils/experiments.py

@@ -7,6 +7,7 @@ import functools
 import hashlib
 import inspect
 import itertools
+import json
 import logging
 import math
 import os
@@ -24,7 +25,7 @@ import warnings
 from abc import ABC, abstractmethod
 from pathlib import Path
 from types import TracebackType
-from typing import Any, Iterator, Self, TypeVar, cast
+from typing import Any, Iterator, Mapping, Self, Sequence, TypeVar, cast
 from urllib.parse import urlparse
 
 import git
@@ -116,9 +117,7 @@ class StateTimer:
 
     def log_dict(self) -> dict[str, int | float | tuple[int | float, bool]]:
         return {
-            "steps": (self.step_timer.steps, True),
             "steps/second": self.step_timer.steps_per_second,
-            "samples": (self.sample_timer.steps, True),
             "samples/second": (self.sample_timer.steps_per_second, True),
             "dt": self.iter_timer.iter_seconds,
         }
@@ -204,8 +203,8 @@ class MinGradScaleError(TrainingFinishedError):
 
 
 def diff_configs(
-    first: ListConfig | DictConfig,
-    second: ListConfig | DictConfig,
+    first: Mapping | Sequence,
+    second: Mapping | Sequence,
     prefix: str | None = None,
 ) -> tuple[list[str], list[str]]:
     """Returns the difference between two configs.
@@ -232,7 +231,7 @@ def diff_configs(
 
     any_config = (ListConfig, DictConfig)
 
-    if isinstance(first, DictConfig) and isinstance(second, DictConfig):
+    if isinstance(first, Mapping) and isinstance(second, Mapping):
         first_keys, second_keys = cast(set[str], set(first.keys())), cast(set[str], set(second.keys()))
 
         # Gets the new keys in each config.
@@ -242,11 +241,12 @@ def diff_configs(
         # Gets the new sub-keys in each config.
         for key in first_keys.intersection(second_keys):
             sub_prefix = key if prefix is None else f"{prefix}.{key}"
-            if OmegaConf.is_missing(first, key) or OmegaConf.is_missing(second, key):
-                if not OmegaConf.is_missing(first, key):
-                    new_first += [get_diff_string(sub_prefix, first[key])]
-                if not OmegaConf.is_missing(second, key):
-                    new_second += [get_diff_string(sub_prefix, second[key])]
+            if isinstance(first, DictConfig) and isinstance(second, DictConfig):
+                if OmegaConf.is_missing(first, key) or OmegaConf.is_missing(second, key):
+                    if not OmegaConf.is_missing(first, key):
+                        new_first += [get_diff_string(sub_prefix, first[key])]
+                    if not OmegaConf.is_missing(second, key):
+                        new_second += [get_diff_string(sub_prefix, second[key])]
             elif isinstance(first[key], any_config) and isinstance(second[key], any_config):
                 sub_new_first, sub_new_second = diff_configs(first[key], second[key], prefix=sub_prefix)
                 new_first, new_second = new_first + sub_new_first, new_second + sub_new_second
@@ -255,7 +255,7 @@ def diff_configs(
                 new_first += [get_diff_string(sub_prefix, first_val)]
                 new_second += [get_diff_string(sub_prefix, second_val)]
 
-    elif isinstance(first, ListConfig) and isinstance(second, ListConfig):
+    elif isinstance(first, Sequence) and isinstance(second, Sequence):
         if len(first) > len(second):
             for i in range(len(second), len(first)):
                 new_first += [get_diff_string(prefix, first[i])]
@@ -470,16 +470,33 @@ def get_command_line_string() -> str:
     return " ".join(sys.argv)
 
 
+def get_environment_variables() -> str:
+    return "\n".join([f"{key}={value}" for key, value in sorted(os.environ.items())])
+
+
 def get_state_file_string(obj: object) -> str:
     return "\n\n".join(
         [
             f"=== Command Line ===\n\n{get_command_line_string()}",
             f"=== Git State ===\n\n{get_git_state(obj)}",
             f"=== Packages ===\n\n{get_packages_with_versions()}",
+            f"=== Environment Variables ===\n\n{get_environment_variables()}",
         ]
     )
 
 
+def get_info_json() -> str:
+    return json.dumps(
+        {
+            "process_id": os.getpid(),
+            "job": {
+                "start_time": datetime.datetime.now().isoformat(),
+            },
+        },
+        indent=2,
+    )
+
+
 def get_training_code(obj: object) -> str:
     """Gets the text from the file containing the provided object.
 

xax/utils/tensorboard.py

@@ -2,11 +2,12 @@
 
 import functools
 import io
+import json
 import os
 import tempfile
 import time
 from pathlib import Path
-from typing import Literal, TypedDict
+from typing import Any, Literal, TypedDict
 
 import numpy as np
 import PIL.Image
@@ -14,9 +15,15 @@ from PIL.Image import Image as PILImage
 from tensorboard.compat.proto.config_pb2 import RunMetadata
 from tensorboard.compat.proto.event_pb2 import Event, TaggedRunMetadata
 from tensorboard.compat.proto.graph_pb2 import GraphDef
-from tensorboard.compat.proto.summary_pb2 import HistogramProto, Summary, SummaryMetadata
+from tensorboard.compat.proto.summary_pb2 import (
+    HistogramProto,
+    Summary,
+    SummaryMetadata,
+)
 from tensorboard.compat.proto.tensor_pb2 import TensorProto
 from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto
+from tensorboard.plugins.mesh import metadata as mesh_metadata
+from tensorboard.plugins.mesh.plugin_data_pb2 import MeshPluginData
 from tensorboard.plugins.text.plugin_data_pb2 import TextPluginData
 from tensorboard.summary.writer.event_file_writer import EventFileWriter
 
@@ -84,6 +91,68 @@ def make_histogram(values: np.ndarray, bins: str | np.ndarray, max_bins: int | N
     )
 
 
+def _get_json_config(config_dict: dict[str, Any] | None) -> str:
+    json_config = "{}"
+    if config_dict is not None:
+        json_config = json.dumps(config_dict, sort_keys=True)
+    return json_config
+
+
+def make_mesh_summary(
+    tag: str,
+    vertices: np.ndarray,
+    colors: np.ndarray | None,
+    faces: np.ndarray | None,
+    config_dict: dict[str, Any] | None,
+    display_name: str | None = None,
+    description: str | None = None,
+) -> Summary:
+    json_config = _get_json_config(config_dict)
+
+    summaries = []
+    tensors = [
+        (vertices, MeshPluginData.VERTEX),
+        (faces, MeshPluginData.FACE),
+        (colors, MeshPluginData.COLOR),
+    ]
+    # Filter out None tensors and explicitly type the list
+    valid_tensors = [(t, content_type) for t, content_type in tensors if t is not None]
+    components = mesh_metadata.get_components_bitmask([content_type for (_, content_type) in valid_tensors])
+
+    for tensor, content_type in valid_tensors:  # Now we know tensor is not None
+        tensor_metadata = mesh_metadata.create_summary_metadata(
+            tag,
+            display_name,
+            content_type,
+            components,
+            tensor.shape,  # Safe now since tensor is not None
+            description,
+            json_config=json_config,
+        )
+
+        tensor_proto = TensorProto(
+            dtype="DT_FLOAT",
+            float_val=tensor.reshape(-1).tolist(),  # Safe now since tensor is not None
+            tensor_shape=TensorShapeProto(
+                dim=[
+                    TensorShapeProto.Dim(size=tensor.shape[0]),  # Safe now since tensor is not None
+                    TensorShapeProto.Dim(size=tensor.shape[1]),
+                    TensorShapeProto.Dim(size=tensor.shape[2]),
+                ]
+            ),
+        )
+
+        tensor_summary = Summary.Value(
+            tag=mesh_metadata.get_instance_name(tag, content_type),
+            tensor=tensor_proto,
+            metadata=tensor_metadata,
+        )
+
+        summaries.append(tensor_summary)
+
+    return Summary(value=summaries)
+
+
 class TensorboardProtobufWriter:
     def __init__(
         self,
@@ -454,6 +523,9 @@ class TensorboardWriter:
         weighted_sum = float((bin_centers * bucket_counts).sum())
         weighted_sum_squares = float((bin_centers**2 * bucket_counts).sum())
 
+        # Convert bin edges to list of floats explicitly
+        bucket_limits: list[float | np.ndarray] = [float(x) for x in bin_edges[1:]]
+
         self.add_histogram_raw(
             tag=tag,
             min=float(bin_edges[0]),
@@ -461,12 +533,28 @@ class TensorboardWriter:
             num=int(total_counts),
             sum=weighted_sum,
             sum_squares=weighted_sum_squares,
-            bucket_limits=bin_edges[1:].tolist(),  # TensorBoard expects right bin edges
+            bucket_limits=bucket_limits,  # Now properly typed
             bucket_counts=bucket_counts.tolist(),
             global_step=global_step,
             walltime=walltime,
         )
 
+    def add_mesh(
+        self,
+        tag: str,
+        vertices: np.ndarray,
+        colors: np.ndarray | None,
+        faces: np.ndarray | None,
+        config_dict: dict[str, Any] | None,
+        global_step: int | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        self.pb_writer.add_summary(
+            make_mesh_summary(tag, vertices, colors, faces, config_dict),
+            global_step=global_step,
+            walltime=walltime,
+        )
+
 
 class TensorboardWriterKwargs(TypedDict):
     max_queue_size: int

{xax-0.1.16.dist-info → xax-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.1.16
+Version: 0.2.1
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -8,14 +8,14 @@ Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: attrs
+Requires-Dist: chex
+Requires-Dist: dpshdl
+Requires-Dist: equinox
+Requires-Dist: importlib-resources
 Requires-Dist: jax
 Requires-Dist: jaxtyping
-Requires-Dist: equinox
 Requires-Dist: optax
-Requires-Dist: dpshdl
-Requires-Dist: chex
-Requires-Dist: importlib-resources
-Requires-Dist: cloudpickle
+Requires-Dist: orbax-checkpoint
 Requires-Dist: pillow
 Requires-Dist: omegaconf
 Requires-Dist: gitpython

{xax-0.1.16.dist-info → xax-0.2.1.dist-info}/RECORD

@@ -1,23 +1,23 @@
-xax/__init__.py,sha256=pSWV5RtPBJynHr7dCqscbnMkETZPUyw8D6MHK4CuS90,14104
+xax/__init__.py,sha256=kd-88OQGnuHb91PXwroAfLb0bMfbe37fXqpECRrjhoU,14182
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
-xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
+xax/requirements.txt,sha256=6qY-84e-sTmlfJNrSjwONQKqzAn5h8G_oGIhnhmfSr4,302
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
-xax/core/state.py,sha256=WwW0qDm-be9MMOT-bGWEFvaWF4iq2FP9xRSn1zq_4A8,2507
+xax/core/state.py,sha256=XejW1tGINYFFcNrscK8eZQsq02J7_RXa461QpmyWuLk,3337
 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=7Yemw3T33QGEP8RkmTkpu6tRVOhut2RUJmttNFfCgFw,5537
 xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
-xax/nn/geom.py,sha256=PN0Ndn575aVtsSfxi67RghHB7luRkqtpS7bPbT1LpLE,5201
+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
 xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
 xax/nn/ssm.py,sha256=8dLAcQ1hBaMT-kkHvwGu_ecxJeTY32WeMYmd4T4KtxA,10745
 xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/base.py,sha256=DqgGIlo5kEWpYix3DdPCEkCgVLUOocjyFr8okaSUq-k,7680
-xax/task/logger.py,sha256=Upx7cCZvaVIs75CHTfIzYmsuaFRsGu0FvziTZuazT4k,37083
+xax/task/base.py,sha256=OnXi2hiKPGwt6ng1dutnoQSiw7lEiWFlC_vx99_JsbQ,7694
+xax/task/logger.py,sha256=peGtfnvnBKr9l6tx1V6XAsvPs0HP6ubV_aE7IJtOMNk,40868
 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
@@ -29,37 +29,37 @@ xax/task/loggers/callback.py,sha256=lyuZX6Bir7xJM07ifdQIl1jlclgkiS82UO9V4y7wgPs,
 xax/task/loggers/json.py,sha256=_tKum6jk_gqVzO-4MqSNXbE-Mmn-yJzkRAT-N1y2zes,4139
 xax/task/loggers/state.py,sha256=6bG-NRsSUzAukYiglCT0oDj8zRMpffH4e1TKWGw1x4k,959
 xax/task/loggers/stdout.py,sha256=oeIgPkj4RyJgBuWaJK9ncLa65iBNJCWXhSF8fx3_54c,6564
-xax/task/loggers/tensorboard.py,sha256=HjR-wiCWe0z3nivRzxEZIltzSzka1828bwxWVmMU5Sk,7718
+xax/task/loggers/tensorboard.py,sha256=KOL9l60tLctX-VAdNwe49H48SAJeGxph3sflJpojA-4,8337
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
 xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
-xax/task/mixins/checkpointing.py,sha256=nRddgtasagf0oTZE9LE5IN5JY7jy4BD_M0rlqYp4sCM,8554
+xax/task/mixins/checkpointing.py,sha256=2nJgqFcV-D8W-4j8TR3PvVh1g5hQUOo-_quKO-XlE4U,11398
 xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 xax/task/mixins/cpu_stats.py,sha256=vAjEc3HpPnl56m7vshYX0dXAHJrB98DzVdsYSRqQllc,9371
-xax/task/mixins/data_loader.py,sha256=WjMWk9uACfBMMClLMcLPkE0WNIvlCZnmqyyqLqJpjX0,6545
+xax/task/mixins/data_loader.py,sha256=Tp7zqPdfH2_JuE6J6EP-fEtCQpq9MjKlGHYK7Zh-goU,6599
 xax/task/mixins/gpu_stats.py,sha256=4HU6teEDlqMitLbSx7fbyL4qBJ0PgGy0Ly_Pzife8yo,8795
 xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,2808
 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=4Xr8b5LFueFh-f3k8MIJMv3M46_Aaf65YwCbjtSBQ-U,26393
+xax/task/mixins/train.py,sha256=v9oi9tNsNBYo-Ne_98nCG9qHX6sxvymHjsRDnL6GL-U,30871
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
-xax/utils/experiments.py,sha256=vm_hWfaty_wEHVdoU2ALiBiGJze7IoDJIfXi6pd_a9I,29360
+xax/utils/experiments.py,sha256=Hzl46_9IH5_9cKzxit-FyVUWBH-_lBs00ZciuIdnWO8,29811
 xax/utils/jax.py,sha256=tC0NNelbrSTzwNGluiwLGKtoHhVpgdzrv-xherB3VtY,4752
 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=VFWhT0MQ99KjQyEYM6NFbqYq4_hOZwB23uhowMB4U34,8754
-xax/utils/tensorboard.py,sha256=21czW8WC2SAmwEhz6RLJc_q5HFvNKM4iR1ZycSO5qPE,17058
+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/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.1.16.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.16.dist-info/METADATA,sha256=gfh7iFi7Wz3fJDf2w1KKs8H0uanhn2HFsR67TvP6uZM,1878
-xax-0.1.16.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.16.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.16.dist-info/RECORD,,
+xax-0.2.1.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.1.dist-info/METADATA,sha256=2pOZLKMIcLoQTM-tRqRvVkF57PZyMoALM87UI5B4dtk,1882
+xax-0.2.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+xax-0.2.1.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.1.dist-info/RECORD,,