xax 0.4.1__tar.gz → 0.4.3__tar.gz

{xax-0.4.1/xax.egg-info → xax-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.4.1
+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.1 → xax-0.4.3}/pyproject.toml

@@ -35,6 +35,7 @@ explicit_package_bases = true
 [[tool.mypy.overrides]]
 
 module = [
+    "chex.*",
     "optax.*",
     "setuptools.*",
     "tensorboard.*",

{xax-0.4.1 → xax-0.4.3}/setup.py

@@ -33,6 +33,7 @@ setup(
     tests_require=requirements_dev,
     extras_require={
         "dev": requirements_dev,
+        "wandb": ["wandb[media]"],
     },
     package_data={
         "xax": [

{xax-0.4.1 → xax-0.4.3}/xax/__init__.py

@@ -12,7 +12,7 @@ and running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.4.1"
+__version__ = "0.4.3"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -171,6 +171,7 @@ __all__ = [
     "uncolored",
     "wrapped",
     "FrozenDict",
+    "freeze_dict",
     "HashableArray",
     "hashable_array",
 ]
@@ -371,6 +372,7 @@ NAME_MAP: dict[str, str] = {
     "uncolored": "utils.text",
     "wrapped": "utils.text",
     "FrozenDict": "utils.types.frozen_dict",
+    "freeze_dict": "utils.types.frozen_dict",
     "HashableArray": "utils.types.hashable_array",
     "hashable_array": "utils.types.hashable_array",
 }
@@ -572,7 +574,7 @@ if IMPORT_ALL or TYPE_CHECKING:
         uncolored,
         wrapped,
     )
-    from xax.utils.types.frozen_dict import FrozenDict
+    from xax.utils.types.frozen_dict import FrozenDict, freeze_dict
     from xax.utils.types.hashable_array import HashableArray, hashable_array
 
 del TYPE_CHECKING, IMPORT_ALL

xax-0.4.3/xax/task/launchers/single_process.py

@@ -0,0 +1,141 @@
+"""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:
+    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()
+
+
+class SingleProcessLauncher(BaseLauncher):
+    def launch(
+        self,
+        task: "type[RunnableMixin[Config]]",
+        *cfgs: RawConfigType,
+        use_cli: bool | list[str] = True,
+    ) -> None:
+        logger = configure_logging()
+        configure_devices(logger)
+        run_single_process_training(task, *cfgs, use_cli=use_cli, logger=logger)

xax-0.4.3/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-0.4.3/xax/task/mixins/logger.py

@@ -0,0 +1,203 @@
+"""Defines a mixin for incorporating some logging functionality."""
+
+import os
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from types import TracebackType
+from typing import Any, Generic, Self, TypeVar
+
+import jax
+
+from xax.core.conf import field
+from xax.core.state import State
+from xax.task.base import BaseConfig, BaseTask
+from xax.task.logger import Logger, LoggerImpl
+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):
+    log_interval_seconds: float = field(
+        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)
+
+
+def get_env_var(name: str, default: bool) -> bool:
+    if name not in os.environ:
+        return default
+    return os.environ[name].strip() == "1"
+
+
+class LoggerMixin(BaseTask[Config], Generic[Config]):
+    logger: Logger
+
+    def __init__(self, config: Config) -> None:
+        super().__init__(config)
+
+        self.logger = Logger()
+
+    def log_directory(self) -> Path | None:
+        return None
+
+    def add_logger(self, *logger: LoggerImpl) -> None:
+        self.logger.add_logger(*logger)
+
+    def set_loggers(self) -> None:
+        self.add_logger(
+            StdoutLogger(
+                log_interval_seconds=self.config.log_interval_seconds,
+            )
+            if is_interactive_session()
+            else JsonLogger(
+                log_interval_seconds=self.config.log_interval_seconds,
+            )
+        )
+
+        # If this is also an ArtifactsMixin, we should default add some
+        # additional loggers which log data to the artifacts directory.
+        if isinstance(self, ArtifactsMixin):
+            self.add_logger(
+                StateLogger(
+                    run_directory=self.exp_dir,
+                ),
+                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)
+
+    def __enter__(self) -> Self:
+        self.logger.__enter__()
+        return self
+
+    def __exit__(self, t: type[BaseException] | None, e: BaseException | None, tr: TracebackType | None) -> None:
+        self.logger.__exit__(t, e, tr)
+        return super().__exit__(t, e, tr)

{xax-0.4.1 → xax-0.4.3}/xax/utils/types/frozen_dict.py

@@ -146,3 +146,7 @@ def unfreeze(x: FrozenDict[K, V] | dict[str, Any]) -> dict[Any, Any]:  # noqa: A
         return ys
     else:
         return x
+
+
+def freeze_dict(x: Mapping[K, V]) -> FrozenDict[K, V]:
+    return FrozenDict(x)

{xax-0.4.1 → xax-0.4.3/xax.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xax
-Version: 0.4.1
+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.1 → xax-0.4.3}/xax.egg-info/SOURCES.txt

@@ -44,6 +44,7 @@ xax/task/loggers/json.py
 xax/task/loggers/state.py
 xax/task/loggers/stdout.py
 xax/task/loggers/tensorboard.py
+xax/task/loggers/wandb.py
 xax/task/mixins/__init__.py
 xax/task/mixins/artifacts.py
 xax/task/mixins/checkpointing.py

{xax-0.4.1 → xax-0.4.3}/xax.egg-info/requires.txt

@@ -23,3 +23,6 @@ pytest
 types-pillow
 types-psutil
 types-requests
+
+[wandb]
+wandb[media]

xax-0.4.1/xax/task/launchers/single_process.py

@@ -1,31 +0,0 @@
-"""Defines a launcher to train a model locally, in a single process."""
-
-from typing import TYPE_CHECKING
-
-from xax.task.base import RawConfigType
-from xax.task.launchers.base import BaseLauncher
-from xax.utils.logging import configure_logging
-
-if TYPE_CHECKING:
-    from xax.task.mixins.runnable import Config, RunnableMixin
-
-
-def run_single_process_training(
-    task: "type[RunnableMixin[Config]]",
-    *cfgs: RawConfigType,
-    use_cli: bool | list[str] = True,
-) -> None:
-    logger = configure_logging()
-    task_obj = task.get_task(*cfgs, use_cli=use_cli)
-    task_obj.add_logger_handlers(logger)
-    task_obj.run()
-
-
-class SingleProcessLauncher(BaseLauncher):
-    def launch(
-        self,
-        task: "type[RunnableMixin[Config]]",
-        *cfgs: RawConfigType,
-        use_cli: bool | list[str] = True,
-    ) -> None:
-        run_single_process_training(task, *cfgs, use_cli=use_cli)

xax-0.4.1/xax/task/mixins/logger.py

@@ -1,92 +0,0 @@
-"""Defines a mixin for incorporating some logging functionality."""
-
-import os
-from dataclasses import dataclass
-from pathlib import Path
-from types import TracebackType
-from typing import Generic, Self, TypeVar
-
-import jax
-
-from xax.core.conf import field
-from xax.core.state import State
-from xax.task.base import BaseConfig, BaseTask
-from xax.task.logger import Logger, LoggerImpl
-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.artifacts import ArtifactsMixin
-from xax.utils.text import is_interactive_session
-
-
-@jax.tree_util.register_dataclass
-@dataclass
-class LoggerConfig(BaseConfig):
-    log_interval_seconds: float = field(
-        value=1.0,
-        help="The interval between successive log lines.",
-    )
-    tensorboard_log_interval_seconds: float = field(
-        value=10.0,
-        help="The interval between successive Tensorboard log lines.",
-    )
-
-
-Config = TypeVar("Config", bound=LoggerConfig)
-
-
-def get_env_var(name: str, default: bool) -> bool:
-    if name not in os.environ:
-        return default
-    return os.environ[name].strip() == "1"
-
-
-class LoggerMixin(BaseTask[Config], Generic[Config]):
-    logger: Logger
-
-    def __init__(self, config: Config) -> None:
-        super().__init__(config)
-
-        self.logger = Logger()
-
-    def log_directory(self) -> Path | None:
-        return None
-
-    def add_logger(self, *logger: LoggerImpl) -> None:
-        self.logger.add_logger(*logger)
-
-    def set_loggers(self) -> None:
-        self.add_logger(
-            StdoutLogger(
-                log_interval_seconds=self.config.log_interval_seconds,
-            )
-            if is_interactive_session()
-            else JsonLogger(
-                log_interval_seconds=self.config.log_interval_seconds,
-            )
-        )
-
-        # If this is also an ArtifactsMixin, we should default add some
-        # additional loggers which log data to the artifacts directory.
-        if isinstance(self, ArtifactsMixin):
-            self.add_logger(
-                StateLogger(
-                    run_directory=self.exp_dir,
-                ),
-                TensorboardLogger(
-                    run_directory=self.exp_dir,
-                    log_interval_seconds=self.config.tensorboard_log_interval_seconds,
-                ),
-            )
-
-    def write_logs(self, state: State) -> None:
-        self.logger.write(state)
-
-    def __enter__(self) -> Self:
-        self.logger.__enter__()
-        return self
-
-    def __exit__(self, t: type[BaseException] | None, e: BaseException | None, tr: TracebackType | None) -> None:
-        self.logger.__exit__(t, e, tr)
-        return super().__exit__(t, e, tr)