xax 0.0.1__py3-none-any.whl → 0.0.5__py3-none-any.whl

xax/__init__.py

@@ -1 +1,256 @@
-__version__ = "0.0.1"
+"""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.5"
+
+# 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",
+    "FourierEmbeddings",
+    "IdentityPositionalEmbeddings",
+    "LearnedPositionalEmbeddings",
+    "RotaryEmbeddings",
+    "SinusoidalEmbeddings",
+    "apply_rotary_embeddings",
+    "cast_embedding_kind",
+    "fourier_embeddings",
+    "get_positional_embeddings",
+    "get_rotary_embeddings",
+    "rotary_embeddings",
+    "BaseLauncher",
+    "CliLauncher",
+    "SingleProcessLauncher",
+    "LogImage",
+    "LogLine",
+    "Logger",
+    "LoggerImpl",
+    "CallbackLogger",
+    "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__ += [
+    "Batch",
+    "CollateMode",
+    "EmbeddingKind",
+    "Output",
+    "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",
+    "FourierEmbeddings": "nn.embeddings",
+    "IdentityPositionalEmbeddings": "nn.embeddings",
+    "LearnedPositionalEmbeddings": "nn.embeddings",
+    "RotaryEmbeddings": "nn.embeddings",
+    "SinusoidalEmbeddings": "nn.embeddings",
+    "apply_rotary_embeddings": "nn.embeddings",
+    "cast_embedding_kind": "nn.embeddings",
+    "fourier_embeddings": "nn.embeddings",
+    "get_positional_embeddings": "nn.embeddings",
+    "get_rotary_embeddings": "nn.embeddings",
+    "rotary_embeddings": "nn.embeddings",
+    "BaseLauncher": "task.launchers.base",
+    "CliLauncher": "task.launchers.cli",
+    "SingleProcessLauncher": "task.launchers.single_process",
+    "LogImage": "task.logger",
+    "LogLine": "task.logger",
+    "Logger": "task.logger",
+    "LoggerImpl": "task.logger",
+    "CallbackLogger": "task.loggers.callback",
+    "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(
+    {
+        "Batch": "task.mixins.train",
+        "CollateMode": "utils.data.collate",
+        "EmbeddingKind": "nn.embeddings",
+        "Output": "task.mixins.output",
+        "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.nn.embeddings import (
+        EmbeddingKind,
+        FourierEmbeddings,
+        IdentityPositionalEmbeddings,
+        LearnedPositionalEmbeddings,
+        RotaryEmbeddings,
+        SinusoidalEmbeddings,
+        apply_rotary_embeddings,
+        cast_embedding_kind,
+        fourier_embeddings,
+        get_positional_embeddings,
+        get_rotary_embeddings,
+        rotary_embeddings,
+    )
+    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 Logger, LoggerImpl, LogImage, LogLine
+    from xax.task.loggers.callback import CallbackLogger
+    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.mixins.train import Batch, Output
+    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/core/conf.py

@@ -0,0 +1,193 @@
+"""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")
+    max_workers: int = field(32, help="Maximum number of workers 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/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})