xax 0.1.14__py3-none-any.whl → 0.1.16__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.14"
+__version__ = "0.1.16"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -40,6 +40,7 @@ __all__ = [
     "load_eqx_mlp",
     "make_eqx_mlp",
     "save_eqx",
+    "cubic_bezier_interpolation",
     "euler_to_quat",
     "get_projected_gravity_vector_from_quat",
     "quat_to_euler",
@@ -201,6 +202,7 @@ NAME_MAP: dict[str, str] = {
     "load_eqx_mlp": "nn.equinox",
     "make_eqx_mlp": "nn.equinox",
     "save_eqx": "nn.equinox",
+    "cubic_bezier_interpolation": "nn.geom",
     "euler_to_quat": "nn.geom",
     "get_projected_gravity_vector_from_quat": "nn.geom",
     "quat_to_euler": "nn.geom",
@@ -363,6 +365,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         save_eqx,
     )
     from xax.nn.geom import (
+        cubic_bezier_interpolation,
         euler_to_quat,
         get_projected_gravity_vector_from_quat,
         quat_to_euler,

xax/nn/geom.py

@@ -1,10 +1,10 @@
 """Defines geometry functions."""
 
-import jax
 from jax import numpy as jnp
+from jaxtyping import Array
 
 
-def quat_to_euler(quat_4: jax.Array, eps: float = 1e-6) -> jax.Array:
+def quat_to_euler(quat_4: Array, eps: float = 1e-6) -> Array:
     """Normalizes and converts a quaternion (w, x, y, z) to roll, pitch, yaw.
 
     Args:
@@ -40,7 +40,7 @@ def quat_to_euler(quat_4: jax.Array, eps: float = 1e-6) -> jax.Array:
     return jnp.concatenate([roll, pitch, yaw], axis=-1)
 
 
-def euler_to_quat(euler_3: jax.Array) -> jax.Array:
+def euler_to_quat(euler_3: Array) -> Array:
     """Converts roll, pitch, yaw angles to a quaternion (w, x, y, z).
 
     Args:
@@ -75,7 +75,7 @@ def euler_to_quat(euler_3: jax.Array) -> jax.Array:
     return quat
 
 
-def get_projected_gravity_vector_from_quat(quat: jax.Array, eps: float = 1e-6) -> jax.Array:
+def get_projected_gravity_vector_from_quat(quat: Array, eps: float = 1e-6) -> Array:
     """Calculates the gravity vector projected onto the local frame given a quaternion orientation.
 
     Args:
@@ -101,7 +101,7 @@ def get_projected_gravity_vector_from_quat(quat: jax.Array, eps: float = 1e-6) -
     return jnp.concatenate([gx, gy, -gz], axis=-1)
 
 
-def rotate_vector_by_quat(vector: jax.Array, quat: jax.Array, eps: float = 1e-6) -> jax.Array:
+def rotate_vector_by_quat(vector: Array, quat: Array, eps: float = 1e-6) -> Array:
     """Rotates a vector by a quaternion.
 
     Args:
@@ -156,3 +156,24 @@ def rotate_vector_by_quat(vector: jax.Array, quat: jax.Array, eps: float = 1e-6)
     )
 
     return jnp.concatenate([xx, yy, zz], axis=-1)
+
+
+def cubic_bezier_interpolation(y_start: Array, y_end: Array, x: Array) -> Array:
+    """Cubic bezier interpolation.
+
+    This is a cubic bezier curve that starts at y_start and ends at y_end,
+    and is controlled by the parameter x. The curve is defined by the following formula:
+
+    y(x) = y_start + (y_end - y_start) * (x**3 + 3 * (x**2 * (1 - x)))
+
+    Args:
+        y_start: The start value, shape (*).
+        y_end: The end value, shape (*).
+        x: The interpolation parameter, shape (*).
+
+    Returns:
+        The interpolated value, shape (*).
+    """
+    y_diff = y_end - y_start
+    bezier = x**3 + 3 * (x**2 * (1 - x))
+    return y_start + y_diff * bezier

xax/task/logger.py

@@ -205,6 +205,12 @@ def as_numpy(array: Array) -> np.ndarray:
     return np.array(array)
 
 
+@dataclass(kw_only=True)
+class LogString:
+    value: str
+    secondary: bool
+
+
 @dataclass(kw_only=True)
 class LogImage:
     image: PILImage
@@ -223,6 +229,12 @@ class LogVideo:
     fps: int
 
 
+@dataclass(kw_only=True)
+class LogScalar:
+    value: Number
+    secondary: bool
+
+
 @dataclass(kw_only=True)
 class LogDistribution:
     mean: Number
@@ -243,10 +255,10 @@ class LogHistogram:
 @dataclass(kw_only=True)
 class LogLine:
     state: State
-    scalars: dict[str, dict[str, Number]]
+    scalars: dict[str, dict[str, LogScalar]]
     distributions: dict[str, dict[str, LogDistribution]]
     histograms: dict[str, dict[str, LogHistogram]]
-    strings: dict[str, dict[str, str]]
+    strings: dict[str, dict[str, LogString]]
     images: dict[str, dict[str, LogImage]]
     videos: dict[str, dict[str, LogVideo]]
 
@@ -515,10 +527,10 @@ class Logger:
     """Defines an intermediate container which holds values to log somewhere else."""
 
     def __init__(self, default_namespace: str = DEFAULT_NAMESPACE) -> None:
-        self.scalars: dict[str, dict[str, Callable[[], Number]]] = defaultdict(dict)
+        self.scalars: dict[str, dict[str, Callable[[], LogScalar]]] = defaultdict(dict)
         self.distributions: dict[str, dict[str, Callable[[], LogDistribution]]] = defaultdict(dict)
         self.histograms: dict[str, dict[str, Callable[[], LogHistogram]]] = defaultdict(dict)
-        self.strings: dict[str, dict[str, Callable[[], str]]] = defaultdict(dict)
+        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.default_namespace = default_namespace
@@ -616,13 +628,23 @@ class Logger:
     def resolve_namespace(self, namespace: str | None = None) -> str:
         return "_".join([self.default_namespace if namespace is None else namespace] + NAMESPACE_STACK)
 
-    def log_scalar(self, key: str, value: Callable[[], Number] | Number, *, namespace: str | None = None) -> None:
+    def log_scalar(
+        self,
+        key: str,
+        value: Callable[[], Number] | Number,
+        *,
+        namespace: str | None = None,
+        secondary: bool = False,
+    ) -> None:
         """Logs a scalar value.
 
         Args:
             key: The key being logged
             value: The scalar value being logged
             namespace: An optional logging namespace
+            secondary: If set, treat this as a secondary value (meaning, it is
+                less important than other values, and some downstream loggers
+                will not display it)
         """
         if not self.active:
             raise RuntimeError("The logger is not active")
@@ -632,11 +654,11 @@ class Logger:
             assert value.ndim == 0, f"Scalar must be a 0D array, got shape {value.shape}"
 
         @functools.lru_cache(maxsize=None)
-        def scalar_future() -> Number:
+        def scalar_future() -> LogScalar:
             with ContextTimer() as timer:
                 value_concrete = value() if callable(value) else value
             logger.debug("Scalar Key: %s, Time: %s", key, timer.elapsed_time)
-            return value_concrete
+            return LogScalar(value=value_concrete, secondary=secondary)
 
         self.scalars[namespace][key] = scalar_future
 
@@ -770,21 +792,31 @@ class Logger:
 
         self.histograms[namespace][key] = histogram_future
 
-    def log_string(self, key: str, value: Callable[[], str] | str, *, namespace: str | None = None) -> None:
+    def log_string(
+        self,
+        key: str,
+        value: Callable[[], str] | str,
+        *,
+        namespace: str | None = None,
+        secondary: bool = False,
+    ) -> None:
         """Logs a string value.
 
         Args:
             key: The key being logged
             value: The string value being logged
             namespace: An optional logging namespace
+            secondary: If set, treat this as a secondary value (meaning, it is
+                less important than other values, and some downstream loggers
+                will not display it)
         """
         if not self.active:
             raise RuntimeError("The logger is not active")
         namespace = self.resolve_namespace(namespace)
 
         @functools.lru_cache(maxsize=None)
-        def value_future() -> str:
-            return value() if callable(value) else value
+        def value_future() -> LogString:
+            return LogString(value=value() if callable(value) else value, secondary=secondary)
 
         self.strings[namespace][key] = value_future
 

xax/task/loggers/json.py

@@ -3,11 +3,19 @@
 import json
 import sys
 from dataclasses import asdict
-from typing import Any, Literal, TextIO
+from typing import Any, Literal, Mapping, TextIO
 
 from jaxtyping import Array
 
-from xax.task.logger import LogError, LoggerImpl, LogLine, LogPing, LogStatus
+from xax.task.logger import (
+    LogError,
+    LoggerImpl,
+    LogLine,
+    LogPing,
+    LogScalar,
+    LogStatus,
+    LogString,
+)
 
 
 def get_json_value(value: Any) -> Any:  # noqa: ANN401
@@ -61,14 +69,14 @@ class JsonLogger(LoggerImpl):
     def get_json(self, line: LogLine) -> str:
         data: dict = {"state": asdict(line.state)}
 
-        def add_logs(log: dict[str, dict[str, Any]], data: dict) -> None:
+        def add_logs(log: Mapping[str, Mapping[str, LogScalar | LogString]], data: dict) -> None:
             for namespace, values in log.items():
                 if self.remove_unicode_from_namespaces:
                     namespace = namespace.encode("ascii", errors="ignore").decode("ascii").strip()
                 if namespace not in data:
                     data[namespace] = {}
                 for k, v in values.items():
-                    data[namespace][k] = get_json_value(v)
+                    data[namespace][k] = get_json_value(v.value)
 
         add_logs(line.scalars, data)
         add_logs(line.strings, data)

xax/task/loggers/stdout.py

@@ -4,11 +4,20 @@ import datetime
 import logging
 import sys
 from collections import deque
-from typing import Any, Deque, TextIO
+from typing import Any, Deque, Mapping, TextIO
 
 from jaxtyping import Array
 
-from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
+from xax.task.logger import (
+    LogError,
+    LogErrorSummary,
+    LoggerImpl,
+    LogLine,
+    LogPing,
+    LogScalar,
+    LogStatus,
+    LogString,
+)
 from xax.utils.text import Color, colored, format_timedelta
 
 
@@ -95,20 +104,17 @@ class StdoutLogger(LoggerImpl):
     def write_log_window(self, line: LogLine) -> None:
         namespace_to_lines: dict[str, dict[str, str]] = {}
 
-        def add_logs(log: dict[str, dict[str, Any]], namespace_to_lines: dict[str, dict[str, str]]) -> None:
+        def add_logs(
+            log: Mapping[str, Mapping[str, LogScalar | LogString]],
+            namespace_to_lines: dict[str, dict[str, str]],
+        ) -> None:
             for namespace, values in log.items():
-                if not self.log_timers and namespace.startswith("⌛"):
-                    continue
-                if not self.log_perf and namespace.startswith("🔧"):
-                    continue
-                if not self.log_optim and namespace.startswith("📉"):
-                    continue
-                if not self.log_fp and namespace.startswith("⚖️"):
-                    continue
-                if namespace not in namespace_to_lines:
-                    namespace_to_lines[namespace] = {}
                 for k, v in values.items():
-                    v_str = as_str(v, self.precision)
+                    if v.secondary:
+                        continue
+                    if namespace not in namespace_to_lines:
+                        namespace_to_lines[namespace] = {}
+                    v_str = as_str(v.value, self.precision)
                     namespace_to_lines[namespace][k] = v_str
 
         add_logs(line.scalars, namespace_to_lines)
@@ -116,9 +122,8 @@ class StdoutLogger(LoggerImpl):
         if not namespace_to_lines:
             return
 
-        self.write_fp.write("\n")
         for namespace, lines in sorted(namespace_to_lines.items()):
-            self.write_fp.write(f"{colored(namespace, 'cyan', bold=True)}\n")
+            self.write_fp.write(f"\n{colored(namespace, 'cyan', bold=True)}\n")
             for k, v in lines.items():
                 self.write_fp.write(f" ↪ {k}: {v}\n")
 

xax/task/loggers/tensorboard.py

@@ -158,7 +158,7 @@ class TensorboardLogger(LoggerImpl):
             for scalar_key, scalar_value in scalars.items():
                 writer.add_scalar(
                     f"{namespace}/{scalar_key}",
-                    as_float(scalar_value),
+                    as_float(scalar_value.value),
                     global_step=line.state.num_steps,
                     walltime=walltime,
                 )
@@ -192,7 +192,7 @@ class TensorboardLogger(LoggerImpl):
             for string_key, string_value in strings.items():
                 writer.add_text(
                     f"{namespace}/{string_key}",
-                    string_value,
+                    string_value.value,
                     global_step=line.state.num_steps,
                     walltime=walltime,
                 )

xax/task/mixins/cpu_stats.py

@@ -248,15 +248,15 @@ class CPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
         stats = monitor.get_if_set() if self.config.cpu_stats.only_log_once else monitor.get()
 
         if stats is not None:
-            self.logger.log_scalar("child_procs", stats.num_child_procs, namespace="🔧 cpu")
-            self.logger.log_scalar("percent", stats.cpu_percent, namespace="🔧 cpu")
-            self.logger.log_scalar("child_percent", stats.child_cpu_percent, namespace="🔧 cpu")
-            self.logger.log_scalar("percent", stats.mem_percent, namespace="🔧 mem")
-            self.logger.log_scalar("shared", stats.mem_shared, namespace="🔧 mem")
-            self.logger.log_scalar("child_percent", stats.child_mem_percent, namespace="🔧 mem")
-            self.logger.log_scalar("rss/cur", stats.mem_rss, namespace="🔧 mem")
-            self.logger.log_scalar("rss/total", stats.mem_rss_total, namespace="🔧 mem")
-            self.logger.log_scalar("vms/cur", stats.mem_vms, namespace="🔧 mem")
-            self.logger.log_scalar("vms/total", stats.mem_vms_total, namespace="🔧 mem")
+            self.logger.log_scalar("child_procs", stats.num_child_procs, namespace="🔧 cpu", secondary=True)
+            self.logger.log_scalar("percent", stats.cpu_percent, namespace="🔧 cpu", secondary=True)
+            self.logger.log_scalar("child_percent", stats.child_cpu_percent, namespace="🔧 cpu", secondary=True)
+            self.logger.log_scalar("percent", stats.mem_percent, namespace="🔧 mem", secondary=True)
+            self.logger.log_scalar("shared", stats.mem_shared, namespace="🔧 mem", secondary=True)
+            self.logger.log_scalar("child_percent", stats.child_mem_percent, namespace="🔧 mem", secondary=True)
+            self.logger.log_scalar("rss/cur", stats.mem_rss, namespace="🔧 mem", secondary=True)
+            self.logger.log_scalar("rss/total", stats.mem_rss_total, namespace="🔧 mem", secondary=True)
+            self.logger.log_scalar("vms/cur", stats.mem_vms, namespace="🔧 mem", secondary=True)
+            self.logger.log_scalar("vms/total", stats.mem_vms_total, namespace="🔧 mem", secondary=True)
 
         return state

xax/task/mixins/gpu_stats.py

@@ -264,8 +264,8 @@ class GPUStatsMixin(ProcessMixin[Config], LoggerMixin[Config], Generic[Config]):
         for gpu_stat in stats.values():
             if gpu_stat is None:
                 continue
-            self.logger.log_scalar(f"mem/{gpu_stat.index}", gpu_stat.memory_used, namespace="🔧 gpu")
-            self.logger.log_scalar(f"temp/{gpu_stat.index}", gpu_stat.temperature, namespace="🔧 gpu")
-            self.logger.log_scalar(f"util/{gpu_stat.index}", gpu_stat.utilization, namespace="🔧 gpu")
+            self.logger.log_scalar(f"mem/{gpu_stat.index}", gpu_stat.memory_used, namespace="🔧 gpu", secondary=True)
+            self.logger.log_scalar(f"temp/{gpu_stat.index}", gpu_stat.temperature, namespace="🔧 gpu", secondary=True)
+            self.logger.log_scalar(f"util/{gpu_stat.index}", gpu_stat.utilization, namespace="🔧 gpu", secondary=True)
 
         return state

xax/task/mixins/train.py

@@ -50,8 +50,7 @@ from xax.utils.experiments import (
     TrainingFinishedError,
     diff_configs,
     get_diff_string,
-    get_git_state,
-    get_packages_with_versions,
+    get_state_file_string,
     get_training_code,
 )
 from xax.utils.jax import jit as xax_jit
@@ -219,7 +218,12 @@ class TrainMixin(
         return state.replace(elapsed_time_s=time.time() - state.start_time_s)
 
     def log_train_step(
-        self, model: PyTree, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State
+        self,
+        model: PyTree,
+        batch: Batch,
+        output: Output,
+        metrics: FrozenDict[str, Array],
+        state: State,
     ) -> None:
         """Override this function to do logging during the training phase.
 
@@ -235,7 +239,12 @@ class TrainMixin(
         """
 
     def log_valid_step(
-        self, model: PyTree, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State
+        self,
+        model: PyTree,
+        batch: Batch,
+        output: Output,
+        metrics: FrozenDict[str, Array],
+        state: State,
     ) -> None:
         """Override this function to do logging during the validation phase.
 
@@ -253,12 +262,20 @@ class TrainMixin(
     def log_state_timers(self, state: State) -> None:
         timer = self.state_timers[state.phase]
         timer.step(state)
-        for ns, d in timer.log_dict().items():
-            for k, v in d.items():
-                self.logger.log_scalar(k, v, namespace=ns)
+        for k, v in timer.log_dict().items():
+            if isinstance(v, tuple):
+                v, secondary = v
+            else:
+                secondary = False
+            self.logger.log_scalar(k, v, namespace="⌛ timers", secondary=secondary)
 
     def log_step(
-        self, model: PyTree, batch: Batch, output: Output, metrics: FrozenDict[str, Array], state: State
+        self,
+        model: PyTree,
+        batch: Batch,
+        output: Output,
+        metrics: FrozenDict[str, Array],
+        state: State,
     ) -> None:
         phase = state.phase
 
@@ -534,9 +551,8 @@ class TrainMixin(
         logger.log(LOG_STATUS, self.task_path)
         logger.log(LOG_STATUS, self.task_name)
         logger.log(LOG_STATUS, "JAX devices: %s", jax.devices())
-        self.logger.log_file("git_state.txt", get_git_state(self))
-        self.logger.log_file("packages.txt", get_packages_with_versions())
-        self.logger.log_file("training_code.txt", get_training_code(self))
+        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))
 
     def model_partition_fn(self, item: Any) -> bool:  # noqa: ANN401

xax/utils/experiments.py

@@ -114,28 +114,15 @@ class StateTimer:
         self.sample_timer.step(state.num_samples if state.phase == "train" else state.num_valid_samples, cur_time)
         self.iter_timer.step(cur_time)
 
-    def log_dict(self) -> dict[str, dict[str, int | float]]:
-        logs: dict[str, dict[str, int | float]] = {}
-
-        # Logs step statistics.
-        logs["⌛ steps"] = {
-            "total": self.step_timer.steps,
-            "per-second": self.step_timer.steps_per_second,
-        }
-
-        # Logs sample statistics.
-        logs["⌛ samples"] = {
-            "total": self.sample_timer.steps,
-            "per-second": self.sample_timer.steps_per_second,
-        }
-
-        # Logs full iteration statistics.
-        logs["⌛ dt"] = {
-            "iter": self.iter_timer.iter_seconds,
+    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,
         }
 
-        return logs
-
 
 class IntervalTicker:
     def __init__(self, interval: float) -> None:
@@ -479,6 +466,20 @@ def get_packages_with_versions() -> str:
     return "\n".join([f"{key}=={version}" for key, version in sorted(packages)])
 
 
+def get_command_line_string() -> str:
+    return " ".join(sys.argv)
+
+
+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()}",
+        ]
+    )
+
+
 def get_training_code(obj: object) -> str:
     """Gets the text from the file containing the provided object.
 

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

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

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

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=D7czvfKKQJlemPuatMPVYbAO4ST3U272QRIyTOru7JI,13989
+xax/__init__.py,sha256=pSWV5RtPBJynHr7dCqscbnMkETZPUyw8D6MHK4CuS90,14104
 xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
 xax/requirements.txt,sha256=9LAEZ5c5gqRSARRVA6xJsVTa4MebPZuC4yOkkwkZJFw,297
@@ -10,14 +10,14 @@ 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=Bj9Z4Y-uoNQuaA_eB_MyG7yImZLuOq8KCLUj1l3daoc,4545
+xax/nn/geom.py,sha256=PN0Ndn575aVtsSfxi67RghHB7luRkqtpS7bPbT1LpLE,5201
 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=1SZjVC6UCtZUoMPcpp3ckotL324QDeYDvHVhf5MHVqg,36271
+xax/task/logger.py,sha256=Upx7cCZvaVIs75CHTfIzYmsuaFRsGu0FvziTZuazT4k,37083
 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
@@ -26,25 +26,25 @@ xax/task/launchers/cli.py,sha256=cK7Nm-3fO-W2gTxpn3FEThsT2NvneS2w0UjA1Nt-84A,140
 xax/task/launchers/single_process.py,sha256=IoML-30g5c526yxkpbWSOtG_KpNQMakT7xujzB1gIAo,846
 xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/task/loggers/callback.py,sha256=lyuZX6Bir7xJM07ifdQIl1jlclgkiS82UO9V4y7wgPs,1582
-xax/task/loggers/json.py,sha256=yXHb1bmfsEnk4p0F1Up1ertWYdcPAFZm25NT8wE3Jb8,4045
+xax/task/loggers/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=BBXqr95gNt5KuCN8XyKnTJF8JdwkR4JgLKrkvcaTBVM,6788
-xax/task/loggers/tensorboard.py,sha256=kI8LvBuBBhPgkP8TeaTQb9SQ0FqaIodwQh2SuWDCnIA,7706
+xax/task/loggers/stdout.py,sha256=oeIgPkj4RyJgBuWaJK9ncLa65iBNJCWXhSF8fx3_54c,6564
+xax/task/loggers/tensorboard.py,sha256=HjR-wiCWe0z3nivRzxEZIltzSzka1828bwxWVmMU5Sk,7718
 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/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,3922
-xax/task/mixins/cpu_stats.py,sha256=C_t71UTrv4LwQzhO5iubsfomj4jYa9bzpE4zBcHdoHM,9211
+xax/task/mixins/cpu_stats.py,sha256=vAjEc3HpPnl56m7vshYX0dXAHJrB98DzVdsYSRqQllc,9371
 xax/task/mixins/data_loader.py,sha256=WjMWk9uACfBMMClLMcLPkE0WNIvlCZnmqyyqLqJpjX0,6545
-xax/task/mixins/gpu_stats.py,sha256=IGPBro9xzSivwD43zM18lWcuei7IhA8LilxSPHqNl4I,8747
+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=aIebtOIvERYofSyqzNGBpNYlNrXweqFUqM9dHiTx3Dc,26253
+xax/task/mixins/train.py,sha256=4Xr8b5LFueFh-f3k8MIJMv3M46_Aaf65YwCbjtSBQ-U,26393
 xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/debugging.py,sha256=OtUdu-3tQsQtik0Q9UM-SNV46IbPjwrAfZcywzoB5d4,1940
-xax/utils/experiments.py,sha256=5CUja1H_cx4dnVqTGQekOpIhqISwHtAgLxZ34GV7cwM,29229
+xax/utils/experiments.py,sha256=vm_hWfaty_wEHVdoU2ALiBiGJze7IoDJIfXi6pd_a9I,29360
 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
@@ -58,8 +58,8 @@ xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,706
 xax/utils/types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xax/utils/types/frozen_dict.py,sha256=ZCMGfSfr2_b2qZbq9ywPD0zej5tpVSId2JftXpwfB5k,4686
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.1.14.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.1.14.dist-info/METADATA,sha256=WbKtAXJUYKHvBrOJPEm_eXF9O9ekc0WdPmsQQCSGG5Q,1878
-xax-0.1.14.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-xax-0.1.14.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.1.14.dist-info/RECORD,,
+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,,