xax 0.0.1__tar.gz → 0.0.3__tar.gz

{xax-0.0.1/xax.egg-info → xax-0.0.3}/PKG-INFO

@@ -1,14 +1,26 @@
 Metadata-Version: 2.1
 Name: xax
-Version: 0.0.1
+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"
@@ -19,3 +31,9 @@ Requires-Dist: ruff; extra == "dev"
 # xax
 
 JAX library for fast experimentation.
+
+## Installation
+
+```bash
+pip install xax
+```

xax-0.0.3/README.md

@@ -0,0 +1,9 @@
+# xax
+
+JAX library for fast experimentation.
+
+## Installation
+
+```bash
+pip install xax
+```

{xax-0.0.1 → xax-0.0.3}/pyproject.toml

@@ -30,11 +30,14 @@ warn_unused_ignores = true
 warn_redundant_casts = true
 
 incremental = true
-namespace_packages = false
+explicit_package_bases = true
 
 [[tool.mypy.overrides]]
 
 module = [
+    "optax.*",
+    "setuptools.*",
+    "tensorboard.*",
     "transformers.*",
 ]
 
@@ -46,33 +49,37 @@ profile = "black"
 
 [tool.ruff]
 
+line-length = 120
+target-version = "py311"
+
+[tool.ruff.lint]
+
 select = ["ANN", "D", "E", "F", "I", "N", "PGH", "PLC", "PLE", "PLR", "PLW", "W"]
 
 ignore = [
     "ANN101", "ANN102", "ANN401",
     "D101", "D102", "D103", "D104", "D105", "D106", "D107",
+    "F722",
     "N812", "N817",
     "PLR0911", "PLR0912", "PLR0913", "PLR0915", "PLR2004",
     "PLW0603", "PLW2901",
 ]
 
-line-length = 120
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
-target-version = "py311"
 
-[tool.ruff.per-file-ignores]
+[tool.ruff.lint.per-file-ignores]
 
 "__init__.py" = ["E402", "F401", "F403", "F811"]
 
-[tool.ruff.mccabe]
+[tool.ruff.lint.mccabe]
 
 max-complexity = 10
 
-[tool.ruff.isort]
+[tool.ruff.lint.isort]
 
 known-first-party = ["xax", "tests"]
 combine-as-imports = true
 
-[tool.ruff.pydocstyle]
+[tool.ruff.lint.pydocstyle]
 
 convention = "google"

{xax-0.0.1 → xax-0.0.3}/setup.cfg

@@ -3,6 +3,8 @@ packages = find:
 
 [options.packages.find]
 exclude = 
+	.vscode
+	.github
 	tests
 
 [egg_info]

{xax-0.0.1 → xax-0.0.3}/setup.py

@@ -35,4 +35,10 @@ setup(
     install_requires=requirements,
     tests_require=requirements_dev,
     extras_require={"dev": requirements_dev},
+    package_data={
+        "xax": [
+            "py.typed",
+            "requirements*.txt",
+        ],
+    },
 )

xax-0.0.3/tests/test_dummy.py

@@ -0,0 +1,5 @@
+"""Runs a dummy placeholder test."""
+
+
+def test_dummy() -> None:
+    assert True

xax-0.0.3/xax/__init__.py

@@ -0,0 +1,214 @@
+"""Defines the top-level xax API.
+
+This package is structured so that all the important stuff can be accessed
+without having to dig around through the internals. This is done by lazily
+importing the module by name.
+
+This file can be maintained by running the update script:
+
+.. code-block:: bash
+
+    python -m scripts.update_api --inplace
+"""
+
+__version__ = "0.0.3"
+
+# This list shouldn't be modified by hand; instead, run the update script.
+__all__ = [
+    "UserConfig",
+    "field",
+    "get_data_dir",
+    "get_pretrained_models_dir",
+    "get_run_dir",
+    "load_user_config",
+    "State",
+    "cast_phase",
+    "BaseLauncher",
+    "CliLauncher",
+    "SingleProcessLauncher",
+    "LogAudio",
+    "LogImage",
+    "LogLine",
+    "LogVideo",
+    "Logger",
+    "LoggerImpl",
+    "JsonLogger",
+    "StateLogger",
+    "StdoutLogger",
+    "TensorboardLogger",
+    "CPUStatsOptions",
+    "DataloaderConfig",
+    "GPUStatsOptions",
+    "Script",
+    "ScriptConfig",
+    "Config",
+    "Task",
+    "collate",
+    "collate_non_null",
+    "BaseFileDownloader",
+    "DataDownloader",
+    "ModelDownloader",
+    "check_md5",
+    "check_sha256",
+    "get_git_state",
+    "get_state_dict_prefix",
+    "get_training_code",
+    "save_config",
+    "ColoredFormatter",
+    "configure_logging",
+    "one_hot",
+    "partial_flatten",
+    "worker_chunk",
+    "TextBlock",
+    "colored",
+    "format_datetime",
+    "format_timedelta",
+    "outlined",
+    "render_text_blocks",
+    "show_error",
+    "show_warning",
+    "uncolored",
+    "wrapped",
+]
+
+__all__ += [
+    "CollateMode",
+    "Phase",
+]
+
+import os
+from typing import TYPE_CHECKING
+
+# If this flag is set, eagerly imports the entire package (not recommended).
+IMPORT_ALL = int(os.environ.get("XAX_IMPORT_ALL", "0")) != 0
+
+del os
+
+# This dictionary is auto-generated and shouldn't be modified by hand; instead,
+# run the update script.
+NAME_MAP: dict[str, str] = {
+    "UserConfig": "core.conf",
+    "field": "core.conf",
+    "get_data_dir": "core.conf",
+    "get_pretrained_models_dir": "core.conf",
+    "get_run_dir": "core.conf",
+    "load_user_config": "core.conf",
+    "State": "core.state",
+    "cast_phase": "core.state",
+    "BaseLauncher": "task.launchers.base",
+    "CliLauncher": "task.launchers.cli",
+    "SingleProcessLauncher": "task.launchers.single_process",
+    "LogAudio": "task.logger",
+    "LogImage": "task.logger",
+    "LogLine": "task.logger",
+    "LogVideo": "task.logger",
+    "Logger": "task.logger",
+    "LoggerImpl": "task.logger",
+    "JsonLogger": "task.loggers.json",
+    "StateLogger": "task.loggers.state",
+    "StdoutLogger": "task.loggers.stdout",
+    "TensorboardLogger": "task.loggers.tensorboard",
+    "CPUStatsOptions": "task.mixins.cpu_stats",
+    "DataLoaderConfig": "task.mixins.data_loader",
+    "GPUStatsOptions": "task.mixins.gpu_stats",
+    "Script": "task.script",
+    "ScriptConfig": "task.script",
+    "Config": "task.task",
+    "Task": "task.task",
+    "collate": "utils.data.collate",
+    "collate_non_null": "utils.data.collate",
+    "BaseFileDownloader": "utils.experiments",
+    "DataDownloader": "utils.experiments",
+    "ModelDownloader": "utils.experiments",
+    "check_md5": "utils.experiments",
+    "check_sha256": "utils.experiments",
+    "get_git_state": "utils.experiments",
+    "get_state_dict_prefix": "utils.experiments",
+    "get_training_code": "utils.experiments",
+    "save_config": "utils.experiments",
+    "ColoredFormatter": "utils.logging",
+    "configure_logging": "utils.logging",
+    "one_hot": "utils.numpy",
+    "partial_flatten": "utils.numpy",
+    "worker_chunk": "utils.numpy",
+    "TextBlock": "utils.text",
+    "colored": "utils.text",
+    "format_datetime": "utils.text",
+    "format_timedelta": "utils.text",
+    "outlined": "utils.text",
+    "render_text_blocks": "utils.text",
+    "show_error": "utils.text",
+    "show_warning": "utils.text",
+    "uncolored": "utils.text",
+    "wrapped": "utils.text",
+}
+
+# Need to manually set some values which can't be auto-generated.
+NAME_MAP.update(
+    {
+        "CollateMode": "utils.data.collate",
+        "Phase": "core.state",
+    },
+)
+
+
+def __getattr__(name: str) -> object:
+    if name not in NAME_MAP:
+        raise AttributeError(f"{__name__} has no attribute {name}")
+
+    module_name = f"xax.{NAME_MAP[name]}"
+    module = __import__(module_name, fromlist=[name])
+    return getattr(module, name)
+
+
+if IMPORT_ALL or TYPE_CHECKING:
+    from xax.core.conf import (
+        UserConfig,
+        field,
+        get_data_dir,
+        get_pretrained_models_dir,
+        get_run_dir,
+        load_user_config,
+    )
+    from xax.core.state import Phase, State, cast_phase
+    from xax.task.launchers.base import BaseLauncher
+    from xax.task.launchers.cli import CliLauncher
+    from xax.task.launchers.single_process import SingleProcessLauncher
+    from xax.task.logger import LogAudio, Logger, LoggerImpl, LogImage, LogLine, LogVideo
+    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.cpu_stats import CPUStatsOptions
+    from xax.task.mixins.data_loader import DataloaderConfig
+    from xax.task.mixins.gpu_stats import GPUStatsOptions
+    from xax.task.script import Script, ScriptConfig
+    from xax.task.task import Config, Task
+    from xax.utils.data.collate import CollateMode, collate, collate_non_null
+    from xax.utils.experiments import (
+        BaseFileDownloader,
+        DataDownloader,
+        ModelDownloader,
+        check_md5,
+        check_sha256,
+        get_git_state,
+        get_state_dict_prefix,
+        get_training_code,
+        save_config,
+    )
+    from xax.utils.logging import ColoredFormatter, configure_logging
+    from xax.utils.numpy import one_hot, partial_flatten, worker_chunk
+    from xax.utils.text import (
+        TextBlock,
+        colored,
+        format_datetime,
+        format_timedelta,
+        outlined,
+        render_text_blocks,
+        show_error,
+        show_warning,
+        uncolored,
+        wrapped,
+    )
+
+del TYPE_CHECKING, IMPORT_ALL

xax-0.0.3/xax/core/conf.py

@@ -0,0 +1,192 @@
+"""Defines base configuration functions and utilities."""
+
+import functools
+import os
+from dataclasses import dataclass, field as field_base
+from pathlib import Path
+from typing import Any, cast
+
+import jax.numpy as jnp
+from omegaconf import II, MISSING, Container as OmegaConfContainer, OmegaConf
+
+from xax.utils.text import show_error
+
+FieldType = Any
+
+
+def field(value: FieldType, **kwargs: str) -> FieldType:
+    """Short-hand function for getting a config field.
+
+    Args:
+        value: The current field's default value.
+        kwargs: Additional metadata fields to supply.
+
+    Returns:
+        The dataclass field.
+    """
+    metadata: dict[str, Any] = {}
+    metadata.update(kwargs)
+
+    if hasattr(value, "__call__"):
+        return field_base(default_factory=value, metadata=metadata)
+    if value.__class__.__hash__ is None:
+        return field_base(default_factory=lambda: value, metadata=metadata)
+    return field_base(default=value, metadata=metadata)
+
+
+def is_missing(cfg: Any, key: str) -> bool:  # noqa: ANN401
+    """Utility function for checking if a config key is missing.
+
+    This is for cases when you are using a raw dataclass rather than an
+    OmegaConf container but want to treat them the same way.
+
+    Args:
+        cfg: The config to check
+        key: The key to check
+
+    Returns:
+        Whether or not the key is missing a value in the config
+    """
+    if isinstance(cfg, OmegaConfContainer):
+        if OmegaConf.is_missing(cfg, key):
+            return True
+        if OmegaConf.is_interpolation(cfg, key):
+            try:
+                getattr(cfg, key)
+                return False
+            except Exception:
+                return True
+    if getattr(cfg, key) is MISSING:
+        return True
+    return False
+
+
+@dataclass
+class Logging:
+    hide_third_party_logs: bool = field(True, help="If set, hide third-party logs")
+    log_level: str = field("INFO", help="The logging level to use")
+
+
+@dataclass
+class Device:
+    cpu: bool = field(True, help="Whether to use the CPU")
+    gpu: bool = field(II("oc.env:USE_GPU,1"), help="Whether to use the GPU")
+    metal: bool = field(II("oc.env:USE_METAL,1"), help="Whether to use the Apple Silicon accelerator")
+    use_fp64: bool = field(False, help="Always use the 64-bit floating point type")
+    use_fp32: bool = field(False, help="Always use the 32-bit floating point type")
+    use_bf16: bool = field(False, help="Always use the 16-bit bfloat type")
+    use_fp16: bool = field(False, help="Always use the 16-bit floating point type")
+
+
+def parse_dtype(cfg: Device) -> jnp.dtype | None:
+    if cfg.use_fp64:
+        return jnp.float64
+    if cfg.use_fp32:
+        return jnp.float32
+    if cfg.use_bf16:
+        return jnp.bfloat16
+    if cfg.use_fp16:
+        return jnp.float16
+    return None
+
+
+@dataclass
+class Triton:
+    use_triton_if_available: bool = field(True, help="Use Triton if available")
+
+
+@dataclass
+class Experiment:
+    default_random_seed: int = field(1337, help="The default random seed to use")
+
+
+@dataclass
+class Directories:
+    run: str = field(II("oc.env:RUN_DIR"), help="The run directory")
+    data: str = field(II("oc.env:DATA_DIR"), help="The data directory")
+    pretrained_models: str = field(II("oc.env:MODEL_DIR"), help="The models directory")
+
+
+@dataclass
+class SlurmPartition:
+    partition: str = field(MISSING, help="The partition name")
+    num_nodes: int = field(1, help="The number of nodes to use")
+
+
+@dataclass
+class Slurm:
+    launch: dict[str, SlurmPartition] = field({}, help="The available launch configurations")
+
+
+@dataclass
+class UserConfig:
+    logging: Logging = field(Logging)
+    device: Device = field(Device)
+    triton: Triton = field(Triton)
+    experiment: Experiment = field(Experiment)
+    directories: Directories = field(Directories)
+    slurm: Slurm = field(Slurm)
+
+
+def user_config_path() -> Path:
+    xaxrc_path_raw = os.environ.get("XAXRC_PATH", "~/.xax.yml")
+    xaxrc_path = Path(xaxrc_path_raw).expanduser()
+    return xaxrc_path
+
+
+@functools.lru_cache(maxsize=None)
+def _load_user_config_cached() -> UserConfig:
+    xaxrc_path = user_config_path()
+    base_cfg = OmegaConf.structured(UserConfig)
+
+    # Writes the config file.
+    if xaxrc_path.exists():
+        cfg = OmegaConf.merge(base_cfg, OmegaConf.load(xaxrc_path))
+    else:
+        show_error(f"No config file was found in {xaxrc_path}; writing one...", important=True)
+        OmegaConf.save(base_cfg, xaxrc_path)
+        cfg = base_cfg
+
+    # Looks in the current directory for a config file.
+    local_cfg_path = Path("xax.yml")
+    if local_cfg_path.exists():
+        cfg = OmegaConf.merge(cfg, OmegaConf.load(local_cfg_path))
+
+    return cast(UserConfig, cfg)
+
+
+def load_user_config() -> UserConfig:
+    """Loads the ``~/.xax.yml`` configuration file.
+
+    Returns:
+        The loaded configuration.
+    """
+    return _load_user_config_cached()
+
+
+def get_run_dir() -> Path | None:
+    config = load_user_config().directories
+    if is_missing(config, "run"):
+        return None
+    (run_dir := Path(config.run)).mkdir(parents=True, exist_ok=True)
+    return run_dir
+
+
+def get_data_dir() -> Path:
+    config = load_user_config().directories
+    if is_missing(config, "data"):
+        raise RuntimeError(
+            "The data directory has not been set! You should set it in your config file "
+            f"in {user_config_path()} or set the DATA_DIR environment variable."
+        )
+    return Path(config.data)
+
+
+def get_pretrained_models_dir() -> Path:
+    config = load_user_config().directories
+    if is_missing(config, "pretrained_models"):
+        raise RuntimeError(
+            "The data directory has not been set! You should set it in your config file "
+            f"in {user_config_path()} or set the MODEL_DIR environment variable."
+        )
+    return Path(config.pretrained_models)

xax-0.0.3/xax/core/state.py

@@ -0,0 +1,81 @@
+"""Defines a dataclass for keeping track of the current training state."""
+
+import time
+from dataclasses import dataclass
+from typing import Literal, TypedDict, cast, get_args
+
+from omegaconf import MISSING
+
+from xax.core.conf import field
+
+Phase = Literal["train", "valid"]
+
+
+def cast_phase(raw_phase: str) -> Phase:
+    args = get_args(Phase)
+    assert raw_phase in args, f"Invalid phase: '{raw_phase}' Valid options are {args}"
+    return cast(Phase, raw_phase)
+
+
+class StateDict(TypedDict, total=False):
+    num_steps: int
+    num_samples: int
+    num_valid_steps: int
+    num_valid_samples: int
+    start_time_s: float
+    elapsed_time_s: float
+    raw_phase: str
+
+
+@dataclass(frozen=True)
+class State:
+    num_steps: int = field(MISSING, help="Number of steps so far")
+    num_samples: int = field(MISSING, help="Number of sample so far")
+    num_valid_steps: int = field(MISSING, help="Number of validation steps so far")
+    num_valid_samples: int = field(MISSING, help="Number of validation samples so far")
+    start_time_s: float = field(MISSING, help="Start time of training")
+    elapsed_time_s: float = field(MISSING, help="Total elapsed time so far")
+    raw_phase: str = field(MISSING, help="Current training phase")
+
+    @property
+    def phase(self) -> Phase:
+        return cast_phase(self.raw_phase)
+
+    @classmethod
+    def init_state(cls) -> "State":
+        return cls(
+            num_steps=0,
+            num_samples=0,
+            num_valid_steps=0,
+            num_valid_samples=0,
+            start_time_s=time.time(),
+            elapsed_time_s=0.0,
+            raw_phase="train",
+        )
+
+    @property
+    def training(self) -> bool:
+        return self.phase == "train"
+
+    def num_phase_steps(self, phase: Phase) -> int:
+        match phase:
+            case "train":
+                return self.num_steps
+            case "valid":
+                return self.num_valid_steps
+            case _:
+                raise ValueError(f"Invalid phase: {phase}")
+
+    def replace(self, values: StateDict) -> "State":
+        return State(
+            num_steps=values.get("num_steps", self.num_steps),
+            num_samples=values.get("num_samples", self.num_samples),
+            num_valid_steps=values.get("num_valid_steps", self.num_valid_steps),
+            num_valid_samples=values.get("num_valid_samples", self.num_valid_samples),
+            start_time_s=values.get("start_time_s", self.start_time_s),
+            elapsed_time_s=values.get("elapsed_time_s", self.elapsed_time_s),
+            raw_phase=values.get("raw_phase", self.raw_phase),
+        )
+
+    def with_phase(self, phase: Phase) -> "State":
+        return self.replace({"raw_phase": phase})

xax-0.0.3/xax/nn/functions.py

@@ -0,0 +1,73 @@
+# mypy: disable-error-code="override"
+"""Defines helper Torch functions."""
+
+import random
+from dataclasses import is_dataclass
+from typing import Any, Callable, Iterable, Mapping, ParamSpec, Sequence, TypeVar
+
+import numpy as np
+from jaxtyping import Array
+
+from xax.core.conf import load_user_config
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+def recursive_apply(item: Any, func: Callable[[Array], Array]) -> Any:  # noqa: ANN401
+    """Applies a function recursively to tensors in an item.
+
+    Args:
+        item: The item to apply the function to
+        func: The function to apply (for the tensor)
+
+    Returns:
+        The same item, with the function applied
+    """
+    if isinstance(item, (str, int, float)):
+        return item
+    if isinstance(item, Array):
+        return func(item)
+    if is_dataclass(item):
+        return item.__class__(**{k: recursive_apply(v, func) for k, v in item.__dict__.items()})
+    if isinstance(item, Mapping):
+        return {k: recursive_apply(v, func) for k, v in item.items()}
+    if isinstance(item, Sequence):
+        return [recursive_apply(i, func) for i in item]
+    return item
+
+
+def recursive_chunk(item: Any, num_chunks: int, dim: int = 0) -> Iterable[Any]:  # noqa: ANN401
+    """Recursively chunk tensors N times.
+
+    Args:
+        item: The item to recursively chunk
+        num_chunks: The number of splits to make
+        dim: The split dimension
+
+    Yields:
+        ``num_chunks`` chunks of items
+    """
+    if isinstance(item, (str, int, float)):
+        yield from (item for _ in range(num_chunks))
+    elif isinstance(item, np.ndarray):
+        yield from np.array_split(item, num_chunks, axis=dim)
+    elif is_dataclass(item):
+        yield from (
+            item.__class__(**{k: i for k, i in zip(item.__dict__, ii)})
+            for ii in zip(*(recursive_chunk(v, num_chunks, dim) for v in item.__dict__.values()))
+        )
+    elif isinstance(item, Mapping):
+        yield from (dict(zip(item, ii)) for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item.values())))
+    elif isinstance(item, Sequence):
+        yield from (list(ii) for ii in zip(*(recursive_chunk(i, num_chunks, dim) for i in item)))
+    else:
+        yield from (item for _ in range(num_chunks))
+
+
+def set_random_seed(seed: int | None = None, offset: int = 0) -> None:
+    if seed is None:
+        seed = load_user_config().experiment.default_random_seed
+    seed += offset
+    random.seed(seed)
+    np.random.seed(seed)