xax 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.2.6"
+__version__ = "0.2.8"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -43,9 +43,12 @@ __all__ = [
     "cubic_bezier_interpolation",
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
+    "normalize",
     "quat_to_euler",
     "quat_to_rotmat",
     "rotate_vector_by_quat",
+    "rotation6d_to_rotation_matrix",
+    "rotation_matrix_to_rotation6d",
     "cross_entropy",
     "cast_norm_type",
     "get_norm",
@@ -57,8 +60,18 @@ __all__ = [
     "BaseLauncher",
     "CliLauncher",
     "SingleProcessLauncher",
+    "LogDistribution",
+    "LogError",
+    "LogErrorSummary",
+    "LogGraph",
+    "LogHistogram",
     "LogImage",
     "LogLine",
+    "LogMesh",
+    "LogPing",
+    "LogScalar",
+    "LogStatus",
+    "LogVideo",
     "Logger",
     "LoggerImpl",
     "CallbackLogger",
@@ -66,6 +79,7 @@ __all__ = [
     "StateLogger",
     "StdoutLogger",
     "TensorboardLogger",
+    "load_ckpt",
     "CPUStatsOptions",
     "DataloaderConfig",
     "GPUStatsOptions",
@@ -115,6 +129,7 @@ __all__ = [
     "compute_nan_ratio",
     "flatten_array",
     "flatten_pytree",
+    "get_pytree_param_count",
     "pytree_has_nans",
     "reshuffle_pytree",
     "reshuffle_pytree_along_dims",
@@ -207,9 +222,12 @@ NAME_MAP: dict[str, str] = {
     "cubic_bezier_interpolation": "nn.geom",
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
+    "normalize": "nn.geom",
     "quat_to_euler": "nn.geom",
     "quat_to_rotmat": "nn.geom",
     "rotate_vector_by_quat": "nn.geom",
+    "rotation6d_to_rotation_matrix": "nn.geom",
+    "rotation_matrix_to_rotation6d": "nn.geom",
     "cross_entropy": "nn.losses",
     "cast_norm_type": "nn.norm",
     "get_norm": "nn.norm",
@@ -221,8 +239,18 @@ NAME_MAP: dict[str, str] = {
     "BaseLauncher": "task.launchers.base",
     "CliLauncher": "task.launchers.cli",
     "SingleProcessLauncher": "task.launchers.single_process",
+    "LogDistribution": "task.logger",
+    "LogError": "task.logger",
+    "LogErrorSummary": "task.logger",
+    "LogGraph": "task.logger",
+    "LogHistogram": "task.logger",
     "LogImage": "task.logger",
     "LogLine": "task.logger",
+    "LogMesh": "task.logger",
+    "LogPing": "task.logger",
+    "LogScalar": "task.logger",
+    "LogStatus": "task.logger",
+    "LogVideo": "task.logger",
     "Logger": "task.logger",
     "LoggerImpl": "task.logger",
     "CallbackLogger": "task.loggers.callback",
@@ -230,6 +258,7 @@ NAME_MAP: dict[str, str] = {
     "StateLogger": "task.loggers.state",
     "StdoutLogger": "task.loggers.stdout",
     "TensorboardLogger": "task.loggers.tensorboard",
+    "load_ckpt": "task.mixins.checkpointing",
     "CPUStatsOptions": "task.mixins.cpu_stats",
     "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
@@ -279,6 +308,7 @@ NAME_MAP: dict[str, str] = {
     "compute_nan_ratio": "utils.pytree",
     "flatten_array": "utils.pytree",
     "flatten_pytree": "utils.pytree",
+    "get_param_count": "utils.pytree",
     "pytree_has_nans": "utils.pytree",
     "reshuffle_pytree": "utils.pytree",
     "reshuffle_pytree_along_dims": "utils.pytree",
@@ -372,9 +402,12 @@ if IMPORT_ALL or TYPE_CHECKING:
         cubic_bezier_interpolation,
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
+        normalize,
         quat_to_euler,
         quat_to_rotmat,
         rotate_vector_by_quat,
+        rotation6d_to_rotation_matrix,
+        rotation_matrix_to_rotation6d,
     )
     from xax.nn.losses import cross_entropy
     from xax.nn.norm import NormType, cast_norm_type, get_norm
@@ -384,12 +417,28 @@ if IMPORT_ALL or TYPE_CHECKING:
     from xax.task.launchers.base import BaseLauncher
     from xax.task.launchers.cli import CliLauncher
     from xax.task.launchers.single_process import SingleProcessLauncher
-    from xax.task.logger import Logger, LoggerImpl, LogImage, LogLine
+    from xax.task.logger import (
+        LogDistribution,
+        LogError,
+        LogErrorSummary,
+        Logger,
+        LoggerImpl,
+        LogGraph,
+        LogHistogram,
+        LogImage,
+        LogLine,
+        LogMesh,
+        LogPing,
+        LogScalar,
+        LogStatus,
+        LogVideo,
+    )
     from xax.task.loggers.callback import CallbackLogger
     from xax.task.loggers.json import JsonLogger
     from xax.task.loggers.state import StateLogger
     from xax.task.loggers.stdout import StdoutLogger
     from xax.task.loggers.tensorboard import TensorboardLogger
+    from xax.task.mixins.checkpointing import load_ckpt
     from xax.task.mixins.cpu_stats import CPUStatsOptions
     from xax.task.mixins.data_loader import DataloaderConfig
     from xax.task.mixins.gpu_stats import GPUStatsOptions
@@ -439,6 +488,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         compute_nan_ratio,
         flatten_array,
         flatten_pytree,
+        get_pytree_param_count,
         pytree_has_nans,
         reshuffle_pytree,
         reshuffle_pytree_along_dims,

xax/core/conf.py

@@ -26,7 +26,7 @@ def field(value: FieldType, **kwargs: str) -> FieldType:
     metadata: dict[str, Any] = {}
     metadata.update(kwargs)
 
-    if hasattr(value, "__call__"):
+    if hasattr(value, "__call__"):  # noqa: B004
         return field_base(default_factory=value, metadata=metadata)
     if value.__class__.__hash__ is None:
         return field_base(default_factory=lambda: value, metadata=metadata)

xax/nn/equinox.py

@@ -68,8 +68,8 @@ def _infer_activation(activation: ActivationFunction) -> Callable:
         return lambda x: x
     try:
         return getattr(jax.nn, activation)
-    except AttributeError:
-        raise ValueError(f"Activation function `{activation}` not found in `jax.nn`")
+    except AttributeError as err:
+        raise ValueError(f"Activation function `{activation}` not found in `jax.nn`") from err
 
 
 def make_eqx_mlp(hyperparams: MLPHyperParams, *, key: PRNGKeyArray) -> eqx.nn.MLP:
@@ -100,7 +100,7 @@ def make_eqx_mlp(hyperparams: MLPHyperParams, *, key: PRNGKeyArray) -> eqx.nn.ML
 def export_eqx_mlp(
     model: eqx.nn.MLP,
     output_path: str | Path,
-    dtype: jax.numpy.dtype = eqx._misc.default_floating_dtype(),
+    dtype: jax.numpy.dtype | None = None,
 ) -> None:
     """Serialize an Equinox MLP to a .eqx file.
 
@@ -109,6 +109,9 @@ def export_eqx_mlp(
         output_path: The path to save the exported model.
         dtype: The dtype of the model.
     """
+    if dtype is None:
+        dtype = eqx._misc.default_floating_dtype()
+
     activation = model.activation.__name__
     final_activation = model.final_activation.__name__
 

xax/nn/functions.py

@@ -1,5 +1,5 @@
 # mypy: disable-error-code="override"
-"""Defines helper Torch functions."""
+"""Defines helper Jax functions."""
 
 import random
 from dataclasses import is_dataclass
@@ -58,13 +58,16 @@ def recursive_chunk(item: Any, num_chunks: int, dim: int = 0) -> Iterable[Any]:
         yield from np.array_split(item, num_chunks, axis=dim)
     elif is_dataclass(item):
         yield from (
-            item.__class__(**{k: i for k, i in zip(item.__dict__, ii)})
-            for ii in zip(*(recursive_chunk(v, num_chunks, dim) for v in item.__dict__.values()))
+            item.__class__(**{k: i for k, i in zip(item.__dict__, ii, strict=True)})
+            for ii in zip(*(recursive_chunk(v, num_chunks, dim) for v in item.__dict__.values()), strict=False)
         )
     elif isinstance(item, Mapping):
-        yield from (dict(zip(item, ii)) for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item.values())))
+        yield from (
+            dict(zip(item, ii, strict=False))
+            for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item.values()), strict=False)
+        )
     elif isinstance(item, Sequence):
-        yield from (list(ii) for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item)))
+        yield from (list(ii) for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item), strict=False))
     else:
         yield from (item for _ in range(num_chunks))
 

xax/nn/geom.py

@@ -1,5 +1,6 @@
 """Defines geometry functions."""
 
+import chex
 from jax import numpy as jnp
 from jaxtyping import Array
 
@@ -211,3 +212,51 @@ def quat_to_rotmat(quat: Array, eps: float = 1e-6) -> Array:
         ],
         axis=-2,
     )
+
+
+def normalize(v: jnp.ndarray, axis: int = -1, eps: float = 1e-8) -> jnp.ndarray:
+    norm = jnp.linalg.norm(v, axis=axis, keepdims=True)
+    return v / jnp.clip(norm, a_min=eps)
+
+
+def rotation6d_to_rotation_matrix(r6d: jnp.ndarray) -> jnp.ndarray:
+    """Convert 6D rotation representation to rotation matrix.
+
+    From https://arxiv.org/pdf/1812.07035, Appendix B
+
+    Args:
+        r6d: The 6D rotation representation, shape (*, 6).
+
+    Returns:
+        The rotation matrix, shape (*, 3, 3).
+    """
+    chex.assert_shape(r6d, (..., 6))
+    shape = r6d.shape
+    flat = r6d.reshape(-1, 6)
+    a_1 = flat[:, 0:3]
+    a_2 = flat[:, 3:6]
+
+    b_1 = normalize(a_1, axis=-1)
+
+    # Reordered Gram-Schmidt orthonormalization.
+    b_3 = normalize(jnp.cross(b_1, a_2), axis=-1)
+    b_2 = jnp.cross(b_3, b_1)
+
+    rotation_matrix = jnp.stack([b_1, b_2, b_3], axis=-1)
+    return rotation_matrix.reshape(shape[:-1] + (3, 3))
+
+
+def rotation_matrix_to_rotation6d(rotation_matrix: jnp.ndarray) -> jnp.ndarray:
+    """Convert rotation matrix to 6D rotation representation.
+
+    Args:
+        rotation_matrix: The rotation matrix, shape (*, 3, 3).
+
+    Returns:
+        The 6D rotation representation, shape (*, 6).
+    """
+    chex.assert_shape(rotation_matrix, (..., 3, 3))
+    shape = rotation_matrix.shape
+    # Simply concatenate a1 and a2 from SO(3)
+    r6d = jnp.concatenate([rotation_matrix[..., 0], rotation_matrix[..., 1]], axis=-1)
+    return r6d.reshape(shape[:-2] + (6,))

xax/task/base.py

@@ -184,8 +184,8 @@ class BaseTask(Generic[Config]):
 
             # Attempts to load any paths as configs.
             is_path = [Path(arg).is_file() or (task_path / arg).is_file() for arg in args]
-            paths = [arg for arg, is_path in zip(args, is_path) if is_path]
-            non_paths = [arg for arg, is_path in zip(args, is_path) if not is_path]
+            paths = [arg for arg, is_path in zip(args, is_path, strict=True) if is_path]
+            non_paths = [arg for arg, is_path in zip(args, is_path, strict=True) if not is_path]
             if paths:
                 cfg = OmegaConf.merge(cfg, *(get_config(path, task_path) for path in paths))
             cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(non_paths))

xax/task/logger.py

@@ -462,11 +462,11 @@ class LoggerImpl(ABC):
 
         self.tickers = {phase: IntervalTicker(log_interval_seconds) for phase in get_args(Phase)}
 
-    def start(self) -> None:
-        pass
+    @abstractmethod
+    def start(self) -> None: ...
 
-    def stop(self) -> None:
-        pass
+    @abstractmethod
+    def stop(self) -> None: ...
 
     @abstractmethod
     def write(self, line: LogLine) -> None:
@@ -476,6 +476,7 @@ class LoggerImpl(ABC):
             line: The line to write.
         """
 
+    @abstractmethod
     def write_error_summary(self, error_summary: LogErrorSummary) -> None:
         """Handles writing an error summary.
 
@@ -483,6 +484,7 @@ class LoggerImpl(ABC):
             error_summary: The error summary to write.
         """
 
+    @abstractmethod
     def write_error(self, error: LogError) -> None:
         """Handles writing an error line.
 
@@ -490,6 +492,7 @@ class LoggerImpl(ABC):
             error: The error information to write.
         """
 
+    @abstractmethod
     def write_status(self, status: LogStatus) -> None:
         """Handles writing a status line.
 
@@ -497,6 +500,7 @@ class LoggerImpl(ABC):
             status: The status to write.
         """
 
+    @abstractmethod
     def write_ping(self, ping: LogPing) -> None:
         """Handles writing a ping line.
 
@@ -504,6 +508,7 @@ class LoggerImpl(ABC):
             ping: The ping to write.
         """
 
+    @abstractmethod
     def log_file(self, name: str, contents: str) -> None:
         """Logs a large text file.
 
@@ -621,7 +626,7 @@ class Logger:
             return
         line = self.pack(state)
         self.clear()
-        for lg in (lg for lg, should_log in zip(self.loggers, should_log) if should_log):
+        for lg in (lg for lg, should_log in zip(self.loggers, should_log, strict=False) if should_log):
             lg.write(line)
 
     def write_error_summary(self, error_summary: str) -> None:
@@ -1045,7 +1050,7 @@ class Logger:
                         line_spacing=line_spacing,
                         centered=centered,
                     )
-                    for img, label in zip(images, labels)
+                    for img, label in zip(images, labels, strict=True)
                 ]
                 tiled = tile_images([img.image for img in labeled], sep)
 

xax/task/loggers/callback.py

@@ -25,6 +25,12 @@ class CallbackLogger(LoggerImpl):
         self.ping_callback = ping_callback
         self.file_callback = file_callback
 
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
     def write(self, line: LogLine) -> None:
         self.callback(line)
 

xax/task/loggers/json.py

@@ -2,13 +2,13 @@
 
 import json
 import sys
-from dataclasses import asdict
 from typing import Any, Literal, Mapping, TextIO
 
 from jaxtyping import Array
 
 from xax.task.logger import (
     LogError,
+    LogErrorSummary,
     LoggerImpl,
     LogLine,
     LogPing,
@@ -58,6 +58,12 @@ class JsonLogger(LoggerImpl):
         self.line_sep = line_sep
         self.remove_unicode_from_namespaces = remove_unicode_from_namespaces
 
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
     @property
     def fp(self) -> TextIO:
         return self.log_stream
@@ -67,7 +73,7 @@ class JsonLogger(LoggerImpl):
         return self.err_log_stream
 
     def get_json(self, line: LogLine) -> str:
-        data: dict = {"state": asdict(line.state)}
+        data: dict = {"state": line.state.to_dict()}
 
         def add_logs(log: Mapping[str, Mapping[str, LogScalar | LogString]], data: dict) -> None:
             for namespace, values in log.items():
@@ -88,6 +94,12 @@ class JsonLogger(LoggerImpl):
         if self.flush_immediately:
             self.fp.flush()
 
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        pass
+
+    def log_file(self, name: str, contents: str) -> None:
+        pass
+
     def write_error(self, error: LogError) -> None:
         self.err_fp.write(error.message)
         if error.location is not None:

xax/task/loggers/state.py

@@ -3,7 +3,14 @@
 from pathlib import Path
 from typing import Literal
 
-from xax.task.logger import LoggerImpl, LogLine
+from xax.task.logger import (
+    LogError,
+    LogErrorSummary,
+    LoggerImpl,
+    LogLine,
+    LogPing,
+    LogStatus,
+)
 
 
 class StateLogger(LoggerImpl):
@@ -30,3 +37,21 @@ class StateLogger(LoggerImpl):
 
     def write(self, line: LogLine) -> None:
         pass
+
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        pass
+
+    def write_error(self, error: LogError) -> None:
+        pass
+
+    def write_status(self, status: LogStatus) -> None:
+        pass
+
+    def write_ping(self, ping: LogPing) -> None:
+        pass

xax/task/loggers/stdout.py

@@ -79,11 +79,13 @@ class StdoutLogger(LoggerImpl):
         self.error_summary: tuple[str, datetime.datetime] | None = None
 
     def start(self) -> None:
-        return super().start()
+        pass
 
     def stop(self) -> None:
         self.write_queues()
-        return super().stop()
+
+    def log_file(self, name: str, contents: str) -> None:
+        pass
 
     def write_separator(self) -> None:
         self.write_fp.write("\033[2J\033[H")

xax/task/loggers/tensorboard.py

@@ -12,7 +12,7 @@ from typing import TypeVar
 
 from xax.core.state import Phase
 from xax.nn.parallel import is_master
-from xax.task.logger import LoggerImpl, LogLine
+from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
 from xax.utils.jax import as_float
 from xax.utils.logging import LOG_STATUS, port_is_busy
 from xax.utils.tensorboard import TensorboardWriter, TensorboardWriters
@@ -236,3 +236,21 @@ class TensorboardLogger(LoggerImpl):
         for name, contents in self.files.items():
             writer.add_text(name, contents)
         self.files.clear()
+
+    def start(self) -> None:
+        pass
+
+    def stop(self) -> None:
+        pass
+
+    def write_error(self, error: LogError) -> None:
+        pass
+
+    def write_error_summary(self, error_summary: LogErrorSummary) -> None:
+        pass
+
+    def write_ping(self, ping: LogPing) -> None:
+        pass
+
+    def write_status(self, status: LogStatus) -> None:
+        pass

xax/task/mixins/artifacts.py

@@ -31,11 +31,13 @@ Config = TypeVar("Config", bound=ArtifactsConfig)
 
 class ArtifactsMixin(BaseTask[Config]):
     _exp_dir: Path | None
+    _stage_dir: Path | None
 
     def __init__(self, config: Config) -> None:
         super().__init__(config)
 
         self._exp_dir = None
+        self._stage_dir = None
 
     @functools.cached_property
     def run_dir(self) -> Path:
@@ -75,15 +77,16 @@ class ArtifactsMixin(BaseTask[Config]):
         logger.log(LOG_STATUS, self._exp_dir)
         return self._exp_dir
 
-    @functools.lru_cache(maxsize=None)
     def stage_environment(self) -> Path | None:
-        stage_dir = (self.exp_dir / "code").resolve()
-        try:
-            stage_environment(self, stage_dir)
-        except Exception:
-            logger.exception("Failed to stage environment!")
-            return None
-        return stage_dir
+        if self._stage_dir is None:
+            stage_dir = (self.exp_dir / "code").resolve()
+            try:
+                stage_environment(self, stage_dir)
+            except Exception:
+                logger.exception("Failed to stage environment!")
+                return None
+            self._stage_dir = stage_dir
+        return self._stage_dir
 
     def on_training_end(self, state: State) -> State:
         state = super().on_training_end(state)

xax/task/mixins/checkpointing.py

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

xax/task/mixins/train.py

@@ -40,7 +40,7 @@ from xax.core.state import Phase, State
 from xax.nn.functions import set_random_seed
 from xax.nn.parallel import is_master
 from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
-from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin, CheckpointPart
+from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin, CheckpointPart, load_ckpt
 from xax.task.mixins.data_loader import DataloadersConfig, DataloadersMixin
 from xax.task.mixins.logger import LoggerConfig, LoggerMixin
 from xax.task.mixins.runnable import RunnableConfig, RunnableMixin
@@ -57,6 +57,7 @@ from xax.utils.experiments import (
 )
 from xax.utils.jax import jit as xax_jit
 from xax.utils.logging import LOG_PING, LOG_STATUS
+from xax.utils.pytree import get_pytree_param_count
 from xax.utils.text import highlight_exception_message, show_info
 from xax.utils.types.frozen_dict import FrozenDict
 
@@ -360,6 +361,7 @@ class TrainMixin(
         model = self.get_model(key)
         state = State.init_state()
 
+        self.log_model_size(model)
         if not load_optimizer:
             return model, state
 
@@ -450,44 +452,43 @@ class TrainMixin(
         match part:
             case "model_state_config":
                 model_spec = eqx.filter_eval_shape(self.get_model, key)
-                return self.load_ckpt_with_template(path, part="model_state_config", model_template=model_spec)
+                model, state, config = load_ckpt(path, part="model_state_config", model_template=model_spec)
+                config = self.get_config(config, use_cli=False)
+                return model, state, config
 
             case "model":
                 model_spec = eqx.filter_eval_shape(self.get_model, key)
-                return self.load_ckpt_with_template(path, part="model", model_template=model_spec)
-
-            case "config":
-                return self.load_ckpt_with_template(path, part="config")
+                return load_ckpt(path, part="model", model_template=model_spec)
 
             case "opt":
                 optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                return self.load_ckpt_with_template(path, part="opt", optimizer_template=optimizer_spec)
+                return load_ckpt(path, part="opt", optimizer_template=optimizer_spec)
 
             case "opt_state":
                 if model is None:
                     model_spec = eqx.filter_eval_shape(self.get_model, key)
-                    model = self.load_ckpt_with_template(path, part="model", model_template=model_spec)
+                    model = load_ckpt(path, part="model", model_template=model_spec)
                 if optimizer is None:
                     optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                    optimizer = self.load_ckpt_with_template(path, part="opt", optimizer_template=optimizer_spec)
+                    optimizer = load_ckpt(path, part="opt", optimizer_template=optimizer_spec)
                 opt_state_spec = eqx.filter_eval_shape(self.get_initial_opt_state, model, optimizer)
-                return self.load_ckpt_with_template(path, part="opt_state", opt_state_template=opt_state_spec)
+                return load_ckpt(path, part="opt_state", opt_state_template=opt_state_spec)
 
             case "state":
-                return self.load_ckpt_with_template(path, part="state")
+                return load_ckpt(path, part="state")
 
             case "config":
-                return self.load_ckpt_with_template(path, part="config")
+                return self.get_config(load_ckpt(path, part="config"), use_cli=False)
 
             case "all":
                 model_spec = eqx.filter_eval_shape(self.get_model, key)
-                model = self.load_ckpt_with_template(path, part="model", model_template=model_spec)
+                model = load_ckpt(path, part="model", model_template=model_spec)
                 optimizer_spec = eqx.filter_eval_shape(self.get_optimizer)
-                optimizer = self.load_ckpt_with_template(path, part="opt", optimizer_template=optimizer_spec)
+                optimizer = load_ckpt(path, part="opt", optimizer_template=optimizer_spec)
                 opt_state_spec = eqx.filter_eval_shape(self.get_initial_opt_state, model, optimizer)
-                opt_state = self.load_ckpt_with_template(path, part="opt_state", opt_state_template=opt_state_spec)
-                state = self.load_ckpt_with_template(path, part="state")
-                config = self.load_ckpt_with_template(path, part="config")
+                opt_state = load_ckpt(path, part="opt_state", opt_state_template=opt_state_spec)
+                state = load_ckpt(path, part="state")
+                config = self.get_config(load_ckpt(path, part="config"), use_cli=False)
                 return model, optimizer, opt_state, state, config
 
             case _:
@@ -683,6 +684,9 @@ class TrainMixin(
         self.logger.log_file("config.yaml", self.config_str(self.config, use_cli=False))
         self.logger.log_file("info.json", get_info_json())
 
+    def log_model_size(self, model: PyTree) -> None:
+        logger.info("Model size: %s", f"{get_pytree_param_count(model):,}")
+
     def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401
         return eqx.is_inexact_array(item)
 

xax/utils/experiments.py

@@ -749,7 +749,8 @@ class BaseFileDownloader(ABC):
                         f"We detected some HTML elements in the downloaded file. "
                         f"This most likely means that the download triggered an unhandled API response by GDrive. "
                         f"Please report this to torchvision at https://github.com/pytorch/vision/issues including "
-                        f"the response:\n\n{text}"
+                        f"the response:\n\n{text}",
+                        stacklevel=2,
                     )
 
     @classmethod

xax/utils/jaxpr.py

@@ -3,10 +3,10 @@
 from pathlib import Path
 
 import jax
-import jax.core
+import jax.extend.core
 
 
-def save_jaxpr_dot(closed_jaxpr: jax.core.ClosedJaxpr, filename: str | Path) -> None:
+def save_jaxpr_dot(closed_jaxpr: jax.extend.core.ClosedJaxpr, filename: str | Path) -> None:
     """Save the JAXPR to a DOT file.
 
     Example usage:
@@ -30,15 +30,15 @@ def save_jaxpr_dot(closed_jaxpr: jax.core.ClosedJaxpr, filename: str | Path) ->
     with open(filename, "w") as f:
         f.write("digraph Jaxpr {\n")
 
-        var_names: dict[jax.core.Var, str] = {}
+        var_names: dict[jax.extend.core.Var, str] = {}
         var_count = 0
 
-        def get_var_name(var: jax.core.Var) -> str:
+        def get_var_name(var: jax.extend.core.Var) -> str:
             """Get a unique name for a variable."""
             nonlocal var_names, var_count
 
             # Handle Literal objects specially since they're not hashable
-            if isinstance(var, jax.core.Literal):
+            if isinstance(var, jax.extend.core.Literal):
                 # Create a name based on the literal value
                 name = f"lit_{var.val}"
                 return name

xax/utils/pytree.py

@@ -1,6 +1,7 @@
 """Utils for accessing, modifying, and otherwise manipulating pytrees."""
 
 import chex
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 from jax import Array
@@ -57,7 +58,7 @@ def pytree_has_nans(pytree: PyTree) -> Array:
 
 def update_pytree(cond: Array, new: PyTree, original: PyTree) -> PyTree:
     """Update a pytree based on a condition."""
-    # Tricky, need use tree_map because where expects array leafs.
+    # Tricky, need use tree.map because where expects array leafs.
     return jax.tree.map(lambda x, y: jnp.where(cond, x, y), new, original)
 
 
@@ -124,7 +125,7 @@ def reshuffle_pytree(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArr
 def reshuffle_pytree_independently(data: PyTree, batch_shape: tuple[int, ...], rng: PRNGKeyArray) -> PyTree:
     """Reshuffle a rollout array across arbitrary batch dimensions independently of each other."""
     rngs = jax.random.split(rng, len(batch_shape))
-    perms = [jax.random.permutation(rng_i, dim) for rng_i, dim in zip(rngs, batch_shape)]
+    perms = [jax.random.permutation(rng_i, dim) for rng_i, dim in zip(rngs, batch_shape, strict=True)]
     # n-dimensional index grid from permutations
     idx_grids = jnp.meshgrid(*perms, indexing="ij")
 
@@ -236,3 +237,9 @@ def reshuffle_pytree_along_dims(
         return x
 
     return jax.tree.map_with_path(restore_transpose, reshuffled_transposed)
+
+
+def get_pytree_param_count(pytree: PyTree) -> int:
+    """Calculates the total number of parameters in a PyTree."""
+    leaves, _ = jax.tree.flatten(pytree)
+    return sum(x.size for x in leaves if isinstance(x, jnp.ndarray) and eqx.is_inexact_array(x))

xax/utils/text.py

@@ -192,7 +192,7 @@ def render_text_blocks(
         if any(len(row) != len(blocks[0]) for row in blocks):
             raise ValueError("All rows must have the same number of blocks in order to align them")
         widths = [[max(len(line) for line in i.lines) if i.width is None else i.width for i in r] for r in blocks]
-        row_widths = [max(i) for i in zip(*widths)]
+        row_widths = [max(i) for i in zip(*widths, strict=True)]
         for row in blocks:
             for i, block in enumerate(row):
                 block.width = row_widths[i]
@@ -263,7 +263,7 @@ def render_text_blocks(
                             if i >= len(block.lines)
                             else colored(pad(block.lines[i], width, block.center), block.color, bold=block.bold)
                         )
-                        for block, width in zip(row, get_widths(row))
+                        for block, width in zip(row, get_widths(row), strict=True)
                     ]
                 )
                 + " │"

xax/utils/types/frozen_dict.py

@@ -133,12 +133,12 @@ class FrozenDict(Mapping[K, V]):
 
     @classmethod
     def tree_unflatten(cls, keys: tuple[K, ...], values: tuple[Any, ...]) -> "FrozenDict[K, V]":
-        return cls({k: v for k, v in zip(keys, values)}, __unsafe_skip_copy__=True)
+        return cls({k: v for k, v in zip(keys, values, strict=True)}, __unsafe_skip_copy__=True)
 
 
 def unfreeze(x: FrozenDict[K, V] | dict[str, Any]) -> dict[Any, Any]:  # noqa: ANN401
     if isinstance(x, FrozenDict):
-        return jax.tree_util.tree_map(lambda y: y, x._dict)
+        return jax.tree.map(lambda y: y, x._dict)
     elif isinstance(x, dict):
         ys = {}
         for key, value in x.items():

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

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

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

@@ -1,23 +1,23 @@
-xax/__init__.py,sha256=pvOmycOUli9h53XMTlinjJdY4zzw_fvWoF05UQwoItI,14225
+xax/__init__.py,sha256=bNV5gnH1foOgZvLf-Cx_fJrhgwk0YFggqB3kNhbNuxg,15502
 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
 xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
+xax/core/conf.py,sha256=d7Dp_GwKnaxtkztlSrJSM_LR0UYJX_FWTtceIWCBkxc,5138
 xax/core/state.py,sha256=yO25lMoLCUTJlHyLzQxlDbsHC_GZ3HkrKAq5huA7AkU,4552
 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/equinox.py,sha256=JZuSApD4bL0UK5W1nrQtucWYvNWUha07J6LTLk_RX-Y,4910
 xax/nn/export.py,sha256=pRfM2B4hB2EvljysC6AjtgB_7Cn7JtaP3dhYU2stZtY,5545
-xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
-xax/nn/geom.py,sha256=rImNlkHWeoNcY7f84nknizJ6uzsrMhbAtKeb2xAWxNY,6215
+xax/nn/functions.py,sha256=bA5kJYzMtFM8eUqBC086i355zJMAO7k_vPFNSDBI9-s,2814
+xax/nn/geom.py,sha256=B8QE-L-xJWhf9KygTByPUAWe7Clpek4GlTABpsJFMBs,7702
 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=OnXi2hiKPGwt6ng1dutnoQSiw7lEiWFlC_vx99_JsbQ,7694
-xax/task/logger.py,sha256=gE67AaPCfU_1FpxY3t0yNRrIVgtp8Sax9UyOqFYMtzM,40976
+xax/task/base.py,sha256=ymGld72zTSw50PBpX4WwqQDyJgQSagDrLh56aqvLT40,7720
+xax/task/logger.py,sha256=W_BpluYvQai1lh1dDCAj-2_mWUC1buhwJncHygDffjc,41125
 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
@@ -25,14 +25,14 @@ xax/task/launchers/base.py,sha256=8LB_r6YISKu1vq1zk3aVYmiedRr9MxE5IMRocs6unFI,73
 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=_tKum6jk_gqVzO-4MqSNXbE-Mmn-yJzkRAT-N1y2zes,4139
-xax/task/loggers/state.py,sha256=6bG-NRsSUzAukYiglCT0oDj8zRMpffH4e1TKWGw1x4k,959
-xax/task/loggers/stdout.py,sha256=ERLFrYe61hSSztzyxBRseobHQR72YFYjEd2i_hOeJ20,6595
-xax/task/loggers/tensorboard.py,sha256=KFlsK0zD2ubDqAXYL4Ds7NQ9F-Ke-PHwfhLOYsGcbw4,8306
+xax/task/loggers/callback.py,sha256=zQuV1xCvz47Q3UQqP1D5mBhbVzptvmPR_7hX25vqSk0,1667
+xax/task/loggers/json.py,sha256=6A5wL7kspsXnpPhI_vu0scgd2Z2-WLhw4gbBFm7eZMM,4377
+xax/task/loggers/state.py,sha256=0Jy0NYnY4c0qt0LvNlaTaCKOSqk5SCKln5VdyuQGnIc,1407
+xax/task/loggers/stdout.py,sha256=giKSW2R83YkgRefm3BLkE7t8Pbj5Dux4AgsdJxYIbGo,6619
+xax/task/loggers/tensorboard.py,sha256=gkAalLsYPGjZiiMlqvDWIhNpYCfKWNvnPz3brIv3JaQ,8725
 xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
-xax/task/mixins/artifacts.py,sha256=2ezmZGzPGe3nhsd9KRkeHWWXdbT9m7drzimIfw6v1XY,2892
-xax/task/mixins/checkpointing.py,sha256=8Hi-2G0EA5OFRjgiOutlk7HgkD5b-0GHazOAYxnGytM,11409
+xax/task/mixins/artifacts.py,sha256=Ma7fwsp-SA1w6GcuBSskszj5TB83yxYJm4Ns_EnqkI4,3018
+xax/task/mixins/checkpointing.py,sha256=zqospBFnTbGt_iriiduVfXazINPbzWpwmIs91KAniMY,10147
 xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
 xax/task/mixins/cpu_stats.py,sha256=rO_9a82ZdsNec61ya4FpYE-rWqPhpijRSXsOfc6caFA,9595
 xax/task/mixins/data_loader.py,sha256=Tp7zqPdfH2_JuE6J6EP-fEtCQpq9MjKlGHYK7Zh-goU,6599
@@ -41,25 +41,25 @@ xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,280
 xax/task/mixins/process.py,sha256=hqDEsMp_SL6ee97iq26-G0g49OcWZZaX82JD4F22eJU,1781
 xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
-xax/task/mixins/train.py,sha256=X3DIMPqpMRAJvFeeDaw5qFkH70R6yJelBAD8KHHwkUc,31196
+xax/task/mixins/train.py,sha256=VJbtoAjtLWJRRxRoMtJt13HkW_R8fOw7u6oBA5gcurk,31264
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
-xax/utils/experiments.py,sha256=d2H63ECtVOKySMUMrQRqq4kcuZpoXqo-L931usDVAhE,29903
+xax/utils/experiments.py,sha256=bj8BftSHT3fFzfiJ0Co0WvqWo0rUS8kQnQYpVvH8FTM,29942
 xax/utils/jax.py,sha256=KQYUHjN6t6JIWa11aRSO3edcsAgTscw_dExxI6kCd9g,6767
-xax/utils/jaxpr.py,sha256=S80nyEkv188RInzq3kCAdkQCU-bf6s0oPTrCE_LjkRs,2298
+xax/utils/jaxpr.py,sha256=H7pWl48ROXIB1-ZPWYfOn-ou3EBMxYWIwc_A0reJQoo,2333
 xax/utils/logging.py,sha256=GAhTne2rdB4Fa1lzk06DMO15U8MTejn6XTClShC-ZtU,6622
 xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
 xax/utils/profile.py,sha256=-aFdWpgYFvBsBZXSLL4zXrFe3zzsDqzmx4q5f2WOtpQ,1628
-xax/utils/pytree.py,sha256=VFWhT0MQ99KjQyEYM6NFbqYq4_hOZwB23uhowMB4U34,8754
+xax/utils/pytree.py,sha256=Qp4u-jNOo9nd-2vdiuGHXFCTDERo4-zAGwJX_7kG7WM,9045
 xax/utils/tensorboard.py,sha256=P0oIFvX2Qts1H4lkpizhRIpQdD0MNppVMeut0Z94yCs,19878
-xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
+xax/utils/text.py,sha256=xS02aSzdywl3KIaNSpKWcxdd37oYlUJtu9wIjkc1wVc,10654
 xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,7066
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/utils/types/frozen_dict.py,sha256=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
+xax/utils/types/frozen_dict.py,sha256=ebtHENhyUzSjyJTlbMaLtcckQIJ7EtgJiok_40TJZpo,4689
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.2.6.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.2.6.dist-info/METADATA,sha256=g--2dJ3PIYDOWURjNhVmkv5B_VUcI2xY0iBa4oGuuVc,1879
-xax-0.2.6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.2.6.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.2.6.dist-info/RECORD,,
+xax-0.2.8.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.2.8.dist-info/METADATA,sha256=zv6QvF5HWciHILyNE-biwqpNIM6QPROSUaQ6aC-dGRU,1879
+xax-0.2.8.dist-info/WHEEL,sha256=lTU6B6eIfYoiQJTZNc-fyaR6BpL6ehTzU3xGYxn2n8k,91
+xax-0.2.8.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.2.8.dist-info/RECORD,,

{xax-0.2.6.dist-info → xax-0.2.8.dist-info}/WHEEL

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