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

xax/utils/experiments.py

@@ -8,6 +8,7 @@ import hashlib
 import inspect
 import itertools
 import logging
+import math
 import os
 import random
 import re
@@ -30,7 +31,7 @@ import requests
 from jaxtyping import Array
 from omegaconf import MISSING, DictConfig, ListConfig, OmegaConf
 
-from xax.core.conf import get_data_dir, get_pretrained_models_dir
+from xax.core.conf import get_data_dir, get_pretrained_models_dir, load_user_config
 from xax.core.state import State
 from xax.utils.text import colored
 
@@ -665,6 +666,7 @@ class BaseFileDownloader(ABC):
         with requests.Session() as session:
             response = session.get(url, params=params, stream=True)
 
+            token: str | None = None
             for key, value in response.cookies.items():
                 if key.startswith("download_warning"):
                     token = value
@@ -756,3 +758,45 @@ def get_state_dict_prefix(
     if regexp is not None:
         ckpt = {k: v for k, v in ckpt.items() if regexp.match(k)}
     return ckpt
+
+
+def split_n_items_across_workers(n: int, worker_id: int, num_workers: int) -> tuple[int, int]:
+    """Computes offsets for splitting N items across K workers.
+
+    This returns the start and end indices for the items to be processed by the
+    given worker. The end index is exclusive.
+
+    Args:
+        n: The number of items to process.
+        worker_id: The ID of the current worker.
+        num_workers: The total number of workers.
+
+    Returns:
+        The start and end index for the items in the current worker.
+    """
+    assert n >= num_workers, f"n ({n}) must be >= num_workers ({num_workers})"
+    assert 0 <= worker_id < num_workers, f"worker_id ({worker_id}) must be >= 0 and < num_workers ({num_workers})"
+
+    # The number of items to process per worker.
+    items_per_worker = math.ceil(n / num_workers)
+
+    # The start and end indices for the items to process.
+    start = worker_id * items_per_worker
+    end = min(start + items_per_worker, n)
+
+    return start, end
+
+
+def num_workers(default: int) -> int:
+    max_workers = load_user_config().experiment.max_workers
+    if hasattr(os, "sched_getaffinity"):
+        try:
+            return min(len(os.sched_getaffinity(0)), max_workers)
+        except Exception:
+            pass
+    if (cpu_count := os.cpu_count()) is not None:
+        return min(cpu_count, max_workers)
+    return min(default, max_workers)
+
+
+OmegaConf.register_new_resolver("mlfab.num_workers", num_workers, replace=True)

xax/utils/logging.py

@@ -5,6 +5,8 @@ import math
 import socket
 import sys
 
+from omegaconf import OmegaConf
+
 from xax.core.conf import load_user_config
 from xax.utils.text import Color, color_parts, colored
 
@@ -175,6 +177,33 @@ def configure_logging(prefix: str | None = None, *, rank: int | None = None, wor
         logging.getLogger("torch").setLevel(logging.WARNING)
 
 
+def get_unused_port(default: int | None = None) -> int:
+    """Returns an unused port number on the local machine.
+
+    Args:
+        default: A default port to try before trying other ports.
+
+    Returns:
+        A port number which is currently unused
+    """
+    if default is not None:
+        sock = socket.socket()
+        try:
+            sock.bind(("", default))
+            return default
+        except OSError:
+            pass
+        finally:
+            sock.close()
+
+    sock = socket.socket()
+    sock.bind(("", 0))
+    return sock.getsockname()[1]
+
+
+OmegaConf.register_new_resolver("mlfab.unused_port", get_unused_port, replace=True)
+
+
 def port_is_busy(port: int) -> int:
     """Checks whether a port is busy.
 

xax/utils/tensorboard.py

@@ -1,11 +1,15 @@
 """Defines utility functions for interfacing with Tensorboard."""
 
 import functools
+import io
+import os
+import tempfile
 import time
 from pathlib import Path
 from typing import Literal, TypedDict
 
 import numpy as np
+import PIL.Image
 from PIL.Image import Image as PILImage
 from tensorboard.compat.proto.config_pb2 import RunMetadata
 from tensorboard.compat.proto.event_pb2 import Event, TaggedRunMetadata
@@ -117,47 +121,111 @@ class TensorboardWriter:
         value: float,
         global_step: int | None = None,
         walltime: float | None = None,
+        new_style: bool = True,
         double_precision: bool = False,
     ) -> None:
+        if new_style:
+            self.pb_writer.add_summary(
+                Summary(
+                    value=[
+                        Summary.Value(
+                            tag=tag,
+                            tensor=(
+                                TensorProto(double_val=[value], dtype="DT_DOUBLE")
+                                if double_precision
+                                else TensorProto(float_val=[value], dtype="DT_FLOAT")
+                            ),
+                            metadata=SummaryMetadata(
+                                plugin_data=SummaryMetadata.PluginData(
+                                    plugin_name="scalars",
+                                ),
+                            ),
+                        )
+                    ],
+                ),
+                global_step=global_step,
+                walltime=walltime,
+            )
+        else:
+            self.pb_writer.add_summary(
+                Summary(
+                    value=[
+                        Summary.Value(
+                            tag=tag,
+                            simple_value=value,
+                        ),
+                    ],
+                ),
+                global_step=global_step,
+                walltime=walltime,
+            )
+
+    def add_image(
+        self,
+        tag: str,
+        value: PILImage,
+        global_step: int | None = None,
+        walltime: float | None = None,
+    ) -> None:
+        output = io.BytesIO()
+        value.convert("RGB").save(output, format="PNG")
+        image_string = output.getvalue()
+        output.close()
+
         self.pb_writer.add_summary(
             Summary(
                 value=[
                     Summary.Value(
                         tag=tag,
-                        tensor=(
-                            TensorProto(double_val=[value], dtype="DT_DOUBLE")
-                            if double_precision
-                            else TensorProto(float_val=[value], dtype="DT_FLOAT")
+                        image=Summary.Image(
+                            height=value.height,
+                            width=value.width,
+                            colorspace=3,  # RGB
+                            encoded_image_string=image_string,
                         ),
-                        metadata=SummaryMetadata(
-                            plugin_data=SummaryMetadata.PluginData(
-                                plugin_name="scalaras",
-                            ),
-                        ),
-                    )
+                    ),
                 ],
             ),
             global_step=global_step,
             walltime=walltime,
         )
 
-    def add_image(
+    def add_video(
         self,
         tag: str,
-        value: PILImage,
+        value: np.ndarray,
         global_step: int | None = None,
         walltime: float | None = None,
+        fps: int = 30,
     ) -> None:
-        image_data = np.asarray(value.convert("RGB")).transpose(2, 0, 1)  # HWC -> CHW
+        assert value.ndim == 4, "Video must be 4D array (T, H, W, C)"
+        images = [PIL.Image.fromarray(frame) for frame in value]
+
+        # Create temporary file for GIF
+        temp_file = tempfile.NamedTemporaryFile(suffix=".gif", delete=False)
+        try:
+            images[0].save(temp_file.name, save_all=True, append_images=images[1:], duration=int(1000 / fps), loop=0)
+            with open(temp_file.name, "rb") as f:
+                video_string = f.read()
+
+        finally:
+            # Clean up temporary file
+            try:
+                os.remove(temp_file.name)
+            except OSError:
+                pass
+
+        # Add to summary
         self.pb_writer.add_summary(
             Summary(
                 value=[
                     Summary.Value(
                         tag=tag,
-                        tensor=TensorProto(
-                            dtype="DT_FLOAT",
-                            tensor_shape=TensorShapeProto(dim=[TensorShapeProto.Dim(size=s) for s in image_data.shape]),
-                            float_val=image_data.flatten().tolist(),
+                        image=Summary.Image(
+                            height=value.shape[1],
+                            width=value.shape[2],
+                            colorspace=value.shape[3],
+                            encoded_image_string=video_string,
                         ),
                     ),
                 ],
@@ -197,7 +265,6 @@ class TensorboardWriter:
 
 
 class TensorboardWriterKwargs(TypedDict):
-    log_directory: Path
     max_queue_size: int
     flush_seconds: float
     filename_suffix: str
@@ -213,8 +280,9 @@ class TensorboardWriters:
     ) -> None:
         super().__init__()
 
+        self.log_directory = Path(log_directory)
+
         self.kwargs: TensorboardWriterKwargs = {
-            "log_directory": Path(log_directory),
             "max_queue_size": max_queue_size,
             "flush_seconds": flush_seconds,
             "filename_suffix": filename_suffix,
@@ -222,11 +290,11 @@ class TensorboardWriters:
 
     @functools.cached_property
     def train_writer(self) -> TensorboardWriter:
-        return TensorboardWriter(**self.kwargs)
+        return TensorboardWriter(self.log_directory / "train", **self.kwargs)
 
     @functools.cached_property
     def valid_writer(self) -> TensorboardWriter:
-        return TensorboardWriter(**self.kwargs)
+        return TensorboardWriter(self.log_directory / "valid", **self.kwargs)
 
     def writer(self, phase: Phase) -> TensorboardWriter:
         match phase:

xax-0.0.6.dist-info/METADATA

@@ -0,0 +1,50 @@
+Metadata-Version: 2.2
+Name: xax
+Version: 0.0.6
+Summary: The xax project
+Home-page: https://github.com/dpshai/xax
+Author: Benjamin Bolte
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+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: pillow
+Requires-Dist: omegaconf
+Requires-Dist: gitpython
+Requires-Dist: tensorboard
+Requires-Dist: psutil
+Requires-Dist: requests
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: darglint; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: types-pillow; extra == "dev"
+Requires-Dist: types-psutil; extra == "dev"
+Requires-Dist: types-requests; extra == "dev"
+Dynamic: author
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# xax
+
+JAX library for fast experimentation.
+
+## Installation
+
+```bash
+pip install xax
+```

xax-0.0.6.dist-info/RECORD

@@ -0,0 +1,52 @@
+xax/__init__.py,sha256=RTUsDh_R0TFa09q-_U0vd-eCYRC-bCaHqHlayp8U2hU,9736
+xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/requirements-dev.txt,sha256=qkscNkFzWd1S5fump-AKH53rR65v2x5FmboFdy_kKvs,128
+xax/requirements.txt,sha256=NmU9PNJhfLtNqqtWWf8WqMjgbBPCn_yt8oMGAgS7Fno,291
+xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/core/conf.py,sha256=Wuo5WLRWuRTgb8eaihvnG_NZskTu0-P3JkIcl_hKINM,5124
+xax/core/state.py,sha256=y123fL7pMgk25TPG6KN0LRIF_eYnD9eP7OfqtoQJGNE,2178
+xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/nn/embeddings.py,sha256=bQGxBFxkLwi2MQLkRfGaHPH5P_KKB21HdI7VNWTKIOQ,11847
+xax/nn/functions.py,sha256=CI_OmspaQwN9nl4hwefIU3_I7m6gBZwJ9aGK1JGUgr0,2713
+xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
+xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/task/base.py,sha256=LHDmM2c_Ps5cGEzn_QUpmyInD7zJJm3Yt9eSeij2Vus,7297
+xax/task/logger.py,sha256=orN1jmM4SIR2EiYk8bNoJZscmhX1FytADBU6p9qpows,29256
+xax/task/script.py,sha256=4LyXrpj0V36TjAZT4lvQeiOTqa5U2tommHKwgWDCE24,1025
+xax/task/task.py,sha256=UHMpnv__gqMcfbC_L-Hhk-DCnUYlFVsgbNf-v8o8B7U,1424
+xax/task/launchers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/task/launchers/base.py,sha256=8LB_r6YISKu1vq1zk3aVYmiedRr9MxE5IMRocs6unFI,731
+xax/task/launchers/cli.py,sha256=cK7Nm-3fO-W2gTxpn3FEThsT2NvneS2w0UjA1Nt-84A,1402
+xax/task/launchers/single_process.py,sha256=IoML-30g5c526yxkpbWSOtG_KpNQMakT7xujzB1gIAo,846
+xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/task/loggers/callback.py,sha256=lyuZX6Bir7xJM07ifdQIl1jlclgkiS82UO9V4y7wgPs,1582
+xax/task/loggers/json.py,sha256=yXHb1bmfsEnk4p0F1Up1ertWYdcPAFZm25NT8wE3Jb8,4045
+xax/task/loggers/state.py,sha256=6bG-NRsSUzAukYiglCT0oDj8zRMpffH4e1TKWGw1x4k,959
+xax/task/loggers/stdout.py,sha256=nxQXkS9JUR38RKsU9qj0dgePKguK0BFa9nl_BdGO8cE,6758
+xax/task/loggers/tensorboard.py,sha256=FGW96z77oG0Kf3cO6Zznx5U3kJNzPWcuSkpY4RnbFCo,6909
+xax/task/mixins/__init__.py,sha256=D3oU31rB9FeOr9MPLleLt5JFbftUr4sBTwgnwQdc2qA,809
+xax/task/mixins/artifacts.py,sha256=1H7ZbR-KSsXhVtqGVlqMi-TXfn1-dM7YnTCLVuw594s,3835
+xax/task/mixins/checkpointing.py,sha256=AMlobojybvJdDZcNCxm1DHSVC_2Qvnu_MbRcsc_8eoA,8508
+xax/task/mixins/compile.py,sha256=FRsxwLnZjjxpeWJ7Bx_d8XUY50oDoGidgpeRt4ejeQk,3377
+xax/task/mixins/cpu_stats.py,sha256=C_t71UTrv4LwQzhO5iubsfomj4jYa9bzpE4zBcHdoHM,9211
+xax/task/mixins/data_loader.py,sha256=WjMWk9uACfBMMClLMcLPkE0WNIvlCZnmqyyqLqJpjX0,6545
+xax/task/mixins/gpu_stats.py,sha256=IGPBro9xzSivwD43zM18lWcuei7IhA8LilxSPHqNl4I,8747
+xax/task/mixins/logger.py,sha256=CIQ4w4K3FcxN6A9xUfITdVkulSxPa4iaTe6cbs9ruaM,1958
+xax/task/mixins/process.py,sha256=d1opVgvc6bOFXb7R58b07F4P5lbSZIzYaajtE0eBbpw,1477
+xax/task/mixins/runnable.py,sha256=IYIsLd2k09g-_y6o44EhJqT7E6BpsyEMmsyLSuzqjtc,1979
+xax/task/mixins/step_wrapper.py,sha256=DJw42mUGwgKx2tkeqatKR9_F4J8ug4wmxKMeJPmhcVQ,1560
+xax/task/mixins/train.py,sha256=dhGL_IuDaJy39BooYlO7JO-_EotKldtBhBplDGU_AnM,21745
+xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/utils/experiments.py,sha256=qT3H0fyVH8DN417x7T0Xmz4SKoogW81-EHcZfyktFI8,28300
+xax/utils/jax.py,sha256=VzEVB766UyH3_cgN6UP0FkCsDuGlYg5KJj8YJS4yYUk,439
+xax/utils/logging.py,sha256=ST1hp2C2xntVVJBUHwo3YxPK19fBLNvHU2WGO1xqcXA,6418
+xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
+xax/utils/tensorboard.py,sha256=oGq2E3Yr0z2xaACv2UOVt_CHEVc8fBxI8V1M99Fd34E,9742
+xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
+xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xax/utils/data/collate.py,sha256=Rd9vMomr_S_zCa_Hi4dO-8ntzAfVwndIUtuXFA3iNcc,7066
+xax-0.0.6.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.0.6.dist-info/METADATA,sha256=YO2c2PUMWkH1ILfPhFWKK4Sodbo9qUpUOCIkm4aLHfg,1171
+xax-0.0.6.dist-info/WHEEL,sha256=jB7zZ3N9hIM9adW7qlTAyycLYW9npaWKLRzaoVcLKcM,91
+xax-0.0.6.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.0.6.dist-info/RECORD,,

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

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

xax/task/launchers/staged.py

@@ -1,29 +0,0 @@
-"""Defines a base class with utility functions for staged training runs."""
-
-from abc import ABC
-from pathlib import Path
-
-from xax.task.launchers.base import BaseLauncher
-from xax.task.mixins.artifacts import ArtifactsMixin, Config
-
-
-class StagedLauncher(BaseLauncher, ABC):
-    def __init__(self, config_file_name: str = "config.yaml") -> None:
-        super().__init__()
-
-        self.config_file_name = config_file_name
-
-    def get_config_path(self, task: "ArtifactsMixin[Config]", use_cli: bool | list[str] = True) -> Path:
-        config_path = task.exp_dir / self.config_file_name
-        task.config.exp_dir = str(task.exp_dir)
-        with open(config_path, "w", encoding="utf-8") as f:
-            f.write(task.config_str(task.config, use_cli=use_cli))
-        return config_path
-
-    @classmethod
-    def from_components(cls, task_key: str, config_path: Path, use_cli: bool | list[str] = True) -> "ArtifactsMixin":
-        return (
-            ArtifactsMixin.from_task_key(task_key)
-            .get_task(config_path, use_cli=use_cli)
-            .set_exp_dir(config_path.parent)
-        )

xax-0.0.3.dist-info/METADATA

@@ -1,39 +0,0 @@
-Metadata-Version: 2.1
-Name: xax
-Version: 0.0.3
-Summary: The xax project
-Home-page: https://github.com/dpshai/xax
-Author: Benjamin Bolte
-Requires-Python: >=3.11
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: dpshdl
-Requires-Dist: equinox
-Requires-Dist: gitpython
-Requires-Dist: jax
-Requires-Dist: jaxtyping
-Requires-Dist: omegaconf
-Requires-Dist: optax
-Requires-Dist: pillow
-Requires-Dist: psutil
-Requires-Dist: requests
-Requires-Dist: tensorboard
-Requires-Dist: types-pillow
-Requires-Dist: types-psutil
-Requires-Dist: types-requests
-Provides-Extra: dev
-Requires-Dist: black ; extra == 'dev'
-Requires-Dist: darglint ; extra == 'dev'
-Requires-Dist: mypy ; extra == 'dev'
-Requires-Dist: pytest ; extra == 'dev'
-Requires-Dist: ruff ; extra == 'dev'
-
-# xax
-
-JAX library for fast experimentation.
-
-## Installation
-
-```bash
-pip install xax
-```

xax-0.0.3.dist-info/RECORD

@@ -1,49 +0,0 @@
-xax/__init__.py,sha256=awE9-hwGPLSx2ywUAUoeBYYm3ztERRhejiZTM8C2ft4,6160
-xax/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/requirements-dev.txt,sha256=aa9ql8CtmuLTvswvwXYlfAxRL0MEyondkMVcuZuM1A8,56
-xax/requirements.txt,sha256=eJ8pC2D3CrbjcETujBG-j3y9dt6WMGpT-wQQzAUh2-4,161
-xax/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/core/conf.py,sha256=DrLXM0VED5RexuqF8-LOljGrnU5soYFgugIrvx6tGmE,5841
-xax/core/state.py,sha256=7lnVSytuhwPfcobPGdjfQ0QxbLgzWQNipKwXchd58QI,2695
-xax/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/nn/functions.py,sha256=ejIltzHagoiKjMEkfVpcIFj_OS3GkMd5JElWUDCrT6Y,2471
-xax/nn/parallel.py,sha256=fnTiT7MsG7eQrJvqwjIz2Ifo3P27TuxIJzmpGYSa_dQ,4608
-xax/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/base.py,sha256=4oueHzToECJYBY-umnrVqDY9tJ6bT18cP97TtaQukBE,7411
-xax/task/logger.py,sha256=lmZ-cfiBisNskz1JDrQmtydT0uwoAAmr38V8-W8XqiU,29200
-xax/task/script.py,sha256=oBGnScYa_X284fCajabPCcbaSEIqR8nO4d40dvMv3NQ,1011
-xax/task/task.py,sha256=mUY7lEug9dW5yrTyT1V5gro5OosnmilMkkWLC-FbI4E,1266
-xax/task/launchers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/launchers/base.py,sha256=8LB_r6YISKu1vq1zk3aVYmiedRr9MxE5IMRocs6unFI,731
-xax/task/launchers/cli.py,sha256=cK7Nm-3fO-W2gTxpn3FEThsT2NvneS2w0UjA1Nt-84A,1402
-xax/task/launchers/single_process.py,sha256=IoML-30g5c526yxkpbWSOtG_KpNQMakT7xujzB1gIAo,846
-xax/task/launchers/staged.py,sha256=jYeT9u58CN4ldV-ltJiQXQglEWOnEckHWnHYjfJQaoY,1102
-xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/task/loggers/json.py,sha256=yXHb1bmfsEnk4p0F1Up1ertWYdcPAFZm25NT8wE3Jb8,4045
-xax/task/loggers/state.py,sha256=qyb-q8MdagN7BX-DhKucwoc45tIZJrPuvVDVoysTKC4,1576
-xax/task/loggers/stdout.py,sha256=nxQXkS9JUR38RKsU9qj0dgePKguK0BFa9nl_BdGO8cE,6758
-xax/task/loggers/tensorboard.py,sha256=Vt9Vr6i5Bmon7z0vykiQKcx_X9qsYTrRgKf1isLQrEA,7388
-xax/task/mixins/__init__.py,sha256=NB8CVbx-zxMpWDYeQJJCUIrUOfrNeVc8mE59qXayEsc,685
-xax/task/mixins/artifacts.py,sha256=W8Y40aqa7QL8mfM5wMRBszLG8Xpv0VZMT3yRnFLfCOU,3568
-xax/task/mixins/cpu_stats.py,sha256=BeG26vkfy4ePAXZfKfoONyHwpCplK-YrbeeEkVNmF8E,8854
-xax/task/mixins/data_loader.py,sha256=CWoozUdNd5UgJ2wzVktHTDBDplb0TZJAThq53sYrKsM,6160
-xax/task/mixins/gpu_stats.py,sha256=HHwBGr56u5MfV_JZDgDvyxnxNgVSr-yNLLfy2tGzyW8,8410
-xax/task/mixins/logger.py,sha256=lWO3nsEdA92Cd8z_JAiD7VBLgiTr2FS5JK2ZJSajexc,9281
-xax/task/mixins/process.py,sha256=am6HeAI9Mw7FmrpYX8tMfIzgq-WSYZwm5T6qKuUlOsw,1327
-xax/task/mixins/runnable.py,sha256=d5-qyIpmNPtbTzE7qFJGGCPSREEDhX1VApUJPNDWye0,1933
-xax/task/mixins/step_wrapper.py,sha256=Do4eGgZVuqDX9ZGDxQdfn6pRbUnHjQBAkTF0vnNH31E,1472
-xax/task/mixins/train.py,sha256=blNJEjUL6FnvHUIhIa8a5au2TUw3OEv57hygfIGfNmI,18626
-xax/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/utils/experiments.py,sha256=EgTTugyBznGhyft04RIA5tDnDGIrib3y3FKa2J6i4sU,26766
-xax/utils/jax.py,sha256=VzEVB766UyH3_cgN6UP0FkCsDuGlYg5KJj8YJS4yYUk,439
-xax/utils/logging.py,sha256=CIOvoH7iq70IRvysG-14OuHSgdvoIfzttwMWiw2yu3E,5732
-xax/utils/numpy.py,sha256=_jOXVi-d2AtJnRftPkRK5MDMzsU8slgw-Jjv4GRm6ns,1197
-xax/utils/tensorboard.py,sha256=BEa1OmOl_7eiq9yahMI2vd6sc_VEBiZnzkEPBGv3IVs,7654
-xax/utils/text.py,sha256=zo1sAoZe59GkpcpaHBVOQ0OekSMGXvOAyNa3lOJozCY,10628
-xax/utils/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xax/utils/data/collate.py,sha256=fuXTzOxdi-STgtRUU2DyJTF72Yz42kWLr6nlszZs6tM,7003
-xax-0.0.3.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.0.3.dist-info/METADATA,sha256=qLBYRcvPbllBPh-VVLDcpSjaYWQJyQxhwrjlx98DTV0,867
-xax-0.0.3.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-xax-0.0.3.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.0.3.dist-info/RECORD,,