xax 0.4.2__py3-none-any.whl → 0.4.3__py3-none-any.whl

xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.4.2"
+__version__ = "0.4.3"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [

xax/task/launchers/single_process.py

@@ -1,21 +1,129 @@
 """Defines a launcher to train a model locally, in a single process."""
 
+import logging
+import os
+import shutil
+import subprocess
 from typing import TYPE_CHECKING
 
+import jax
+
 from xax.task.base import RawConfigType
 from xax.task.launchers.base import BaseLauncher
+from xax.task.mixins.gpu_stats import get_num_gpus
 from xax.utils.logging import configure_logging
 
 if TYPE_CHECKING:
     from xax.task.mixins.runnable import Config, RunnableMixin
 
 
+def get_gpu_memory_info() -> dict[int, tuple[float, float]]:
+    """Get memory information for all GPUs.
+
+    Returns:
+        Dictionary mapping GPU index to (total_memory_mb, used_memory_mb)
+    """
+    command = "nvidia-smi --query-gpu=index,memory.total,memory.used --format=csv,noheader"
+
+    try:
+        with subprocess.Popen(command.split(), stdout=subprocess.PIPE, universal_newlines=True) as proc:
+            stdout = proc.stdout
+            assert stdout is not None
+
+            gpu_info = {}
+            for line in stdout:
+                line = line.strip()
+                if not line:
+                    continue
+
+                parts = line.split(", ")
+                if len(parts) >= 3:
+                    gpu_id = int(parts[0])
+                    total_mem = float(parts[1].replace(" MiB", ""))
+                    used_mem = float(parts[2].replace(" MiB", ""))
+                    gpu_info[gpu_id] = (total_mem, used_mem)
+
+            return gpu_info
+
+    except Exception as e:
+        logger = configure_logging()
+        logger.warning("Failed to get GPU memory info: %s", e)
+        return {}
+
+
+def select_best_gpu() -> int | None:
+    """Select the GPU with the most available memory.
+
+    Returns:
+        GPU index with most available memory, or None if no GPUs found
+    """
+    gpu_info = get_gpu_memory_info()
+
+    if not gpu_info:
+        return None
+
+    # Find GPU with most available memory
+    best_gpu = None
+    max_available: float = -1.0
+
+    for gpu_id, (total_mem, used_mem) in gpu_info.items():
+        available_mem = total_mem - used_mem
+        if available_mem > max_available:
+            max_available = available_mem
+            best_gpu = gpu_id
+
+    return best_gpu
+
+
+def configure_gpu_devices(logger: logging.Logger | None = None) -> None:
+    if logger is None:
+        logger = configure_logging()
+
+    # If there are multiple devices, choose the one with the most
+    # available memory (i.e., the one which is likely not being used
+    # by other processes) and use only that device.
+    num_gpus = get_num_gpus()
+
+    if num_gpus > 1:
+        logger.info("Multiple GPUs detected (%d), selecting GPU with most available memory", num_gpus)
+
+        best_gpu = select_best_gpu()
+        if best_gpu is not None:
+            logger.info("Selected GPU %d for training", best_gpu)
+
+            # Set CUDA_VISIBLE_DEVICES to only show the selected GPU
+            os.environ["CUDA_VISIBLE_DEVICES"] = str(best_gpu)
+
+            # Configure JAX to use the selected device
+            try:
+                devices = jax.devices("gpu")
+                if devices:
+                    jax.config.update("jax_default_device", devices[0])
+                    logger.info("Configured JAX to use device: %s", devices[0])
+            except Exception as e:
+                logger.warning("Failed to configure JAX device: %s", e)
+        else:
+            logger.warning("Could not determine best GPU, using default device selection")
+    elif num_gpus == 1:
+        logger.info("Single GPU detected, using default device selection")
+
+
+def configure_devices(logger: logging.Logger | None = None) -> None:
+    if logger is None:
+        logger = configure_logging()
+
+    if shutil.which("nvidia-smi") is not None:
+        configure_gpu_devices(logger)
+
+
 def run_single_process_training(
     task: "type[RunnableMixin[Config]]",
     *cfgs: RawConfigType,
     use_cli: bool | list[str] = True,
+    logger: logging.Logger | None = None,
 ) -> None:
-    logger = configure_logging()
+    if logger is None:
+        logger = configure_logging()
     task_obj = task.get_task(*cfgs, use_cli=use_cli)
     task_obj.add_logger_handlers(logger)
     task_obj.run()
@@ -28,4 +136,6 @@ class SingleProcessLauncher(BaseLauncher):
         *cfgs: RawConfigType,
         use_cli: bool | list[str] = True,
     ) -> None:
-        run_single_process_training(task, *cfgs, use_cli=use_cli)
+        logger = configure_logging()
+        configure_devices(logger)
+        run_single_process_training(task, *cfgs, use_cli=use_cli, logger=logger)

xax/task/loggers/wandb.py

@@ -0,0 +1,307 @@
+# mypy: disable-error-code="import-not-found"
+"""Defines a Weights & Biases logger backend."""
+
+import logging
+import os
+from enum import Enum
+from pathlib import Path
+from typing import Any, TypeVar
+
+import numpy as np
+
+from xax.nn.parallel import is_master
+from xax.task.logger import LogError, LogErrorSummary, LoggerImpl, LogLine, LogPing, LogStatus
+from xax.utils.jax import as_float
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+def sanitize_metric_name(name: str) -> str:
+    """Remove 4-byte unicode characters from metric names.
+
+    W&B has issues with 4-byte unicode characters in metric names,
+    so we need to filter them out.
+
+    Args:
+        name: The metric name to sanitize.
+
+    Returns:
+        The sanitized metric name.
+    """
+    # Filter out characters that don't fit in UCS-2 (Basic Multilingual Plane)
+    # These are characters with code points > 0xFFFF (4-byte UTF-8)
+    return "".join(char for char in name if ord(char) <= 0xFFFF)
+
+
+class WandbConfigResumeOption(str, Enum):
+    ALLOW = "allow"
+    NEVER = "never"
+    MUST = "must"
+    AUTO = "auto"
+
+
+class WandbConfigModeOption(str, Enum):
+    ONLINE = "online"
+    OFFLINE = "offline"
+    DISABLED = "disabled"
+    SHARED = "shared"
+
+
+class WandbConfigReinitOption(str, Enum):
+    RETURN_PREVIOUS = "return_previous"
+    FINISH_PREVIOUS = "finish_previous"
+
+
+WandbConfigResume = WandbConfigResumeOption | bool
+WandbConfigMode = WandbConfigModeOption | None
+
+
+class WandbLogger(LoggerImpl):
+    def __init__(
+        self,
+        project: str | None = None,
+        entity: str | None = None,
+        name: str | None = None,
+        run_directory: str | Path | None = None,
+        config: dict[str, Any] | None = None,
+        tags: list[str] | None = None,
+        notes: str | None = None,
+        log_interval_seconds: float = 10.0,
+        reinit: WandbConfigReinitOption = WandbConfigReinitOption.RETURN_PREVIOUS,
+        resume: WandbConfigResume = False,
+        mode: WandbConfigMode = None,
+    ) -> None:
+        """Defines a logger which writes to Weights & Biases.
+
+        Args:
+            project: The name of the W&B project to log to.
+            entity: The W&B entity (team or user) to log to.
+            name: The name of this run.
+            run_directory: The root run directory. If provided, wandb will save
+                files to a subdirectory here.
+            config: Configuration dictionary to log.
+            tags: List of tags for this run.
+            notes: Notes about this run.
+            log_interval_seconds: The interval between successive log lines.
+            reinit: Whether to allow multiple wandb.init() calls in the same process.
+            resume: Whether to resume a previous run. Can be a run ID string.
+            mode: Mode for wandb ("online", "offline", or "disabled").
+        """
+        try:
+            import wandb as _wandb  # noqa: F401,PLC0415
+        except ImportError as e:
+            raise RuntimeError(
+                "WandbLogger requires the 'wandb' package. Install it with: pip install xax[wandb]"
+            ) from e
+
+        self._wandb = _wandb
+
+        super().__init__(log_interval_seconds)
+
+        self.project = project
+        self.entity = entity
+        self.name = name
+        self.config = config
+        self.tags = tags
+        self.notes = notes
+        self.reinit = reinit
+        self.resume: WandbConfigResume = resume
+        self.mode: WandbConfigMode = mode
+
+        # Set wandb directory if run_directory is provided
+        if run_directory is not None:
+            self.wandb_dir = Path(run_directory).expanduser().resolve() / "wandb"
+            self.wandb_dir.mkdir(parents=True, exist_ok=True)
+
+        self._started = False
+
+        # Store pending files to log
+        self.files: dict[str, str] = {}
+
+        self.start()
+
+    def start(self) -> None:
+        """Initialize the W&B run."""
+        if self._started or not is_master():
+            return
+
+        # Set wandb environment variables if needed
+        if self.wandb_dir is not None:
+            os.environ["WANDB_DIR"] = str(self.wandb_dir)
+
+        # Initialize wandb run
+        self.run = self._wandb.init(  # pyright
+            project=self.project,
+            entity=self.entity,
+            name=self.name,
+            config=self.config,
+            tags=self.tags,
+            notes=self.notes,
+            reinit=self.reinit.value,
+            resume=self.resume.value if isinstance(self.resume, WandbConfigResumeOption) else self.resume,
+            mode=self.mode.value if isinstance(self.mode, WandbConfigModeOption) else self.mode,
+        )
+
+        self._started = True
+        logger.info("W&B run initialized: %s", self.run.url if self.run else "No URL available")
+
+    def stop(self) -> None:
+        """Finish the W&B run."""
+        if not self._started or not is_master():
+            return
+
+        if self.run is not None:
+            self.run.finish()
+        self._started = False
+
+    def log_file(self, name: str, contents: str) -> None:
+        """Store a file to be logged with the next write call.
+
+        Args:
+            name: The name of the file.
+            contents: The contents of the file.
+        """
+        if not is_master():
+            return
+        self.files[name] = contents
+
+    def write(self, line: LogLine) -> None:
+        """Writes the current log line to W&B.
+
+        Args:
+            line: The line to write.
+        """
+        if not is_master() or not self._started:
+            return
+
+        # Get step information
+        global_step = line.state.num_steps.item()
+
+        # Dictionary to collect all metrics for this step
+        metrics: dict[str, Any] = {}
+
+        # Log scalars
+        for namespace, scalars in line.scalars.items():
+            for scalar_key, scalar_value in scalars.items():
+                key = sanitize_metric_name(f"{namespace}/{scalar_key}")
+                metrics[key] = as_float(scalar_value.value)
+
+        # Log distributions as custom metrics (mean and std)
+        for namespace, distributions in line.distributions.items():
+            for distribution_key, distribution_value in distributions.items():
+                base_key = sanitize_metric_name(f"{namespace}/{distribution_key}")
+                metrics[f"{base_key}/mean"] = float(distribution_value.mean)
+                metrics[f"{base_key}/std"] = float(distribution_value.std)
+
+        # Log histograms
+        for namespace, histograms in line.histograms.items():
+            for histogram_key, histogram_value in histograms.items():
+                key = sanitize_metric_name(f"{namespace}/{histogram_key}")
+                # Create histogram data for wandb
+                # W&B expects a list of values or a numpy array
+                # We need to reconstruct the data from the histogram bins
+                values = []
+                for i, count in enumerate(histogram_value.bucket_counts):
+                    if count > 0:
+                        # Use the midpoint of each bucket
+                        if i == 0:
+                            val = histogram_value.bucket_limits[0]
+                        else:
+                            val = (histogram_value.bucket_limits[i - 1] + histogram_value.bucket_limits[i]) / 2
+                        values.extend([val] * count)
+
+                if values:
+                    # wandb.Histogram accepts lists directly
+                    metrics[key] = self._wandb.Histogram(values)
+
+        # Log strings as HTML
+        for namespace, strings in line.strings.items():
+            for string_key, string_value in strings.items():
+                key = sanitize_metric_name(f"{namespace}/{string_key}")
+                # For strings, we can log them as HTML or just as text in a table
+                metrics[key] = self._wandb.Html(f"<pre>{string_value.value}</pre>")
+
+        # Log images
+        for namespace, images in line.images.items():
+            for image_key, image_value in images.items():
+                key = sanitize_metric_name(f"{namespace}/{image_key}")
+                # Convert PIL image to wandb.Image
+                metrics[key] = self._wandb.Image(image_value.image)
+
+        # Log videos
+        for namespace, videos in line.videos.items():
+            for video_key, video_value in videos.items():
+                key = sanitize_metric_name(f"{namespace}/{video_key}")
+                # wandb.Video expects shape (time, channels, height, width)
+                # Our format is (T, H, W, C) so we need to transpose to (T, C, H, W)
+                frames = video_value.frames.transpose(0, 3, 1, 2)  # (T, H, W, C) -> (T, C, H, W)
+                metrics[key] = self._wandb.Video(frames, fps=video_value.fps, format="mp4")
+
+        # Log meshes (3D objects)
+        for namespace, meshes in line.meshes.items():
+            for mesh_key, mesh_value in meshes.items():
+                key = sanitize_metric_name(f"{namespace}/{mesh_key}")
+                # W&B Object3D expects vertices and faces in specific format
+                # vertices: (batch_size, num_vertices, 3) or (num_vertices, 3)
+                # faces: (batch_size, num_faces, 3) or (num_faces, 3)
+                vertices = mesh_value.vertices
+
+                # Handle batch dimension - take first batch if present
+                if vertices.ndim == 3:
+                    vertices = vertices[0]
+
+                obj3d_data = {
+                    "type": "lidar/beta",
+                    "vertices": vertices.tolist(),
+                }
+
+                if mesh_value.faces is not None:
+                    faces = mesh_value.faces
+                    if faces.ndim == 3:
+                        faces = faces[0]
+                    obj3d_data["faces"] = faces.tolist()
+
+                if mesh_value.colors is not None:
+                    colors = mesh_value.colors
+                    if colors.ndim == 3:
+                        colors = colors[0]
+                    # Convert colors to 0-1 range if they're in 0-255 range
+                    # The colors are already numpy arrays from LogMesh, converted by as_numpy
+                    if colors.dtype == np.uint8:
+                        colors = colors.astype(np.float32) / 255.0
+                    obj3d_data["colors"] = colors.tolist()
+
+                metrics[key] = self._wandb.Object3D(obj3d_data)
+
+        # Log any pending files as artifacts or text
+        for name, contents in self.files.items():
+            # Log as HTML text
+            key = sanitize_metric_name(name)
+            key = f"{self.run.name}_{key}"
+            is_training_code = "code" in name
+            artifact = self._wandb.Artifact(
+                name=key if not is_training_code else "training_code",
+                type="code" if is_training_code else "unspecified",
+            )
+            with artifact.new_file(name) as f:
+                f.write(contents)
+            artifact.save()
+        self.files.clear()
+
+        # Log all metrics at once
+        if metrics and self.run:
+            self.run.log(metrics, step=global_step)
+
+    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/mixins/logger.py

@@ -2,9 +2,10 @@
 
 import os
 from dataclasses import dataclass
+from enum import Enum
 from pathlib import Path
 from types import TracebackType
-from typing import Generic, Self, TypeVar
+from typing import Any, Generic, Self, TypeVar
 
 import jax
 
@@ -16,10 +17,22 @@ 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.loggers.wandb import (
+    WandbConfigMode,
+    WandbConfigModeOption,
+    WandbConfigReinitOption,
+    WandbConfigResume,
+    WandbLogger,
+)
 from xax.task.mixins.artifacts import ArtifactsMixin
 from xax.utils.text import is_interactive_session
 
 
+class LoggerBackend(str, Enum):
+    TENSORBOARD = "tensorboard"
+    WANDB = "wandb"
+
+
 @jax.tree_util.register_dataclass
 @dataclass
 class LoggerConfig(BaseConfig):
@@ -27,10 +40,50 @@ class LoggerConfig(BaseConfig):
         value=1.0,
         help="The interval between successive log lines.",
     )
+    logger_backend: LoggerBackend = field(
+        value=LoggerBackend.TENSORBOARD,
+        help="The logger backend to use",
+    )
     tensorboard_log_interval_seconds: float = field(
         value=10.0,
         help="The interval between successive Tensorboard log lines.",
     )
+    wandb_project: str | None = field(
+        value=None,
+        help="The name of the W&B project to log to.",
+    )
+    wandb_entity: str | None = field(
+        value=None,
+        help="The W&B entity (team or user) to log to.",
+    )
+    wandb_name: str | None = field(
+        value=None,
+        help="The name of this run in W&B.",
+    )
+    wandb_tags: list[str] | None = field(
+        value=None,
+        help="List of tags for this W&B run.",
+    )
+    wandb_notes: str | None = field(
+        value=None,
+        help="Notes about this W&B run.",
+    )
+    wandb_log_interval_seconds: float = field(
+        value=10.0,
+        help="The interval between successive W&B log lines.",
+    )
+    wandb_mode: WandbConfigMode = field(
+        value=WandbConfigModeOption.ONLINE,
+        help="Mode for wandb (online, offline, or disabled).",
+    )
+    wandb_resume: WandbConfigResume = field(
+        value=False,
+        help="Whether to resume a previous run. Can be a run ID string.",
+    )
+    wandb_reinit: WandbConfigReinitOption = field(
+        value=WandbConfigReinitOption.RETURN_PREVIOUS,
+        help="Whether to allow multiple wandb.init() calls in the same process.",
+    )
 
 
 Config = TypeVar("Config", bound=LoggerConfig)
@@ -74,12 +127,70 @@ class LoggerMixin(BaseTask[Config], Generic[Config]):
                 StateLogger(
                     run_directory=self.exp_dir,
                 ),
-                TensorboardLogger(
-                    run_directory=self.exp_dir,
-                    log_interval_seconds=self.config.tensorboard_log_interval_seconds,
-                ),
+                self._create_logger_backend(),
             )
 
+    def _create_logger_backend(self) -> LoggerImpl:
+        match self.config.logger_backend:
+            case LoggerBackend.TENSORBOARD:
+                return TensorboardLogger(
+                    run_directory=self.exp_dir if isinstance(self, ArtifactsMixin) else "./",
+                    log_interval_seconds=self.config.tensorboard_log_interval_seconds,
+                )
+            case LoggerBackend.WANDB:
+                run_config = {}
+                if hasattr(self.config, "__dict__"):
+                    # Convert config to a serializable dictionary
+                    run_config = self._config_to_dict(self.config)
+
+                return WandbLogger(
+                    project=self.config.wandb_project,
+                    entity=self.config.wandb_entity,
+                    name=self.config.wandb_name,
+                    run_directory=self.exp_dir if isinstance(self, ArtifactsMixin) else None,
+                    config=run_config,
+                    tags=self.config.wandb_tags,
+                    notes=self.config.wandb_notes,
+                    log_interval_seconds=self.config.wandb_log_interval_seconds,
+                    reinit=self.config.wandb_reinit,
+                    resume=self.config.wandb_resume,
+                    mode=self.config.wandb_mode,
+                )
+            case _:
+                # This shouldn't happen, as validation should take care of this
+                raise Exception(f"Invalid logger_backend '{self.config.logger_backend}'")
+
+    def _config_to_dict(self, config: Config) -> dict[str, Any]:
+        """Convert a config object to a dictionary for W&B logging.
+
+        Args:
+            config: The configuration object to convert.
+
+        Returns:
+            A dictionary representation of the config.
+        """
+        if hasattr(config, "__dict__"):
+            result: dict[str, Any] = {}
+            for key, value in config.__dict__.items():
+                if not key.startswith("_"):
+                    # Recursively convert nested configs
+                    if hasattr(value, "__dict__"):
+                        result[key] = self._config_to_dict(value)
+                    elif isinstance(value, (list, tuple)):
+                        # Handle lists/tuples that might contain configs
+                        result[key] = [
+                            self._config_to_dict(item) if hasattr(item, "__dict__") else item for item in value
+                        ]
+                    elif isinstance(value, dict):
+                        # Handle dicts that might contain configs
+                        result[key] = {
+                            k: self._config_to_dict(v) if hasattr(v, "__dict__") else v for k, v in value.items()
+                        }
+                    else:
+                        result[key] = value
+            return result
+        return {}
+
     def write_logs(self, state: State) -> None:
         self.logger.write(state)
 

{xax-0.4.2.dist-info → xax-0.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.4.2
+Version: 0.4.3
 Summary: A library for fast Jax experimentation
 Home-page: https://github.com/kscalelabs/xax
 Author: Benjamin Bolte
@@ -31,6 +31,8 @@ Requires-Dist: pytest; extra == "dev"
 Requires-Dist: types-pillow; extra == "dev"
 Requires-Dist: types-psutil; extra == "dev"
 Requires-Dist: types-requests; extra == "dev"
+Provides-Extra: wandb
+Requires-Dist: wandb[media]; extra == "wandb"
 Dynamic: author
 Dynamic: description
 Dynamic: description-content-type

{xax-0.4.2.dist-info → xax-0.4.3.dist-info}/RECORD

@@ -1,4 +1,4 @@
-xax/__init__.py,sha256=NsrSPe__5x39x_Csd05Sqt4VbUGsIs9JC48h4gB3T1o,17242
+xax/__init__.py,sha256=ckZ0hysU32yz3jOglRyau1iw2Tt3Ct0QP9qffIORZS4,17242
 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
@@ -25,13 +25,14 @@ xax/task/task.py,sha256=Iy02wRUti5lDX1rfDHIgX87dGYeayJxJ9nzJzp_lMq0,1960
 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=wdEUT-B-FE9aemmt1tB_rKKRNy60aiDhslsy2i-ojWo,896
+xax/task/launchers/single_process.py,sha256=D67OtUGLifZa3wxWsB1RN97D-OjrA9gopAVdb-IlRzA,4471
 xax/task/loggers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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=sRyBbeBeVXDTYhPZIKIapW0JEfL9hqqzhNTeIcSd374,8883
+xax/task/loggers/wandb.py,sha256=iLs8VnFPJ_yryzLX5njrGOvZrqNehnVZM09xfUzRMuU,11392
 xax/task/mixins/__init__.py,sha256=wYc4zfutdMyEmzCVV421gSf25ZXW9htNTSY_TW6vL_8,894
 xax/task/mixins/artifacts.py,sha256=UN26TW22ARduO6Bjs0yRu4-V6-Md9MPbXLKDnS28m44,3861
 xax/task/mixins/checkpointing.py,sha256=v50IZ7j58DWmEu-_6Zh_02R5KUVGhrMkg5n-MYM_J4c,11484
@@ -39,7 +40,7 @@ xax/task/mixins/compile.py,sha256=PG5aF3W9v_xGiImHgUJ7gmwuQQoSQWufdpl2N_mlLX0,39
 xax/task/mixins/cpu_stats.py,sha256=rO_9a82ZdsNec61ya4FpYE-rWqPhpijRSXsOfc6caFA,9595
 xax/task/mixins/data_loader.py,sha256=BKfOVWXR70vbyHMFlnlUiQQHXHH5zTj5WtmsymNCFB4,6722
 xax/task/mixins/gpu_stats.py,sha256=USOyhXldxbsrl6eCtoFKTWUm_lfeG0cUCkQNUpXRdtA,8880
-xax/task/mixins/logger.py,sha256=6oXsJJyNUx6YT3q58FVXMZBUpMgjVkGre6BXFN20cVI,2808
+xax/task/mixins/logger.py,sha256=mgJlw8ZYshm8F88jo1RuRNBIvYwYk7CF2o_7V6LPN_Q,7136
 xax/task/mixins/process.py,sha256=hqDEsMp_SL6ee97iq26-G0g49OcWZZaX82JD4F22eJU,1781
 xax/task/mixins/runnable.py,sha256=pcLrYc_TycZUY9zZim05Skc2FWk3IZKFnu6p3UDMonM,1966
 xax/task/mixins/step_wrapper.py,sha256=-Yu5Nft2CRw1JvZt6J_94SM1vqX8fk08IDK95Pmd2ew,1648
@@ -61,9 +62,9 @@ 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=QBdBblKenWffJyADQaFF_0iAEVa9EgMylV63udA2vJ4,4771
 xax/utils/types/hashable_array.py,sha256=l5iIcFmkYzfGeaZmcSoeFkthFASqM8xJYK3AXhZQYwc,992
-xax-0.4.2.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-xax-0.4.2.dist-info/METADATA,sha256=e7es2HRb5RwbM7K1o75P7htxsxVwNdUcXvzTS9jyLI8,1246
-xax-0.4.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-xax-0.4.2.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
-xax-0.4.2.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
-xax-0.4.2.dist-info/RECORD,,
+xax-0.4.3.dist-info/licenses/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+xax-0.4.3.dist-info/METADATA,sha256=kfEAbt8WILGev-MQY5uFmrTrCFI3uVI5R0lnigd2ZXw,1314
+xax-0.4.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+xax-0.4.3.dist-info/entry_points.txt,sha256=uRC6rx5ce0bf-FblJaZSBMxxKFfMyoWTf8OWbBmLSe8,61
+xax-0.4.3.dist-info/top_level.txt,sha256=g4Au_r2XhvZ-lTybviH-Fh9g0zF4DAYHYxPue1-xbs8,4
+xax-0.4.3.dist-info/RECORD,,