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

xax/__init__.py

@@ -11,7 +11,7 @@ This file can be maintained by running the update script:
     python -m scripts.update_api --inplace
 """
 
-__version__ = "0.0.3"
+__version__ = "0.0.6"
 
 # This list shouldn't be modified by hand; instead, run the update script.
 __all__ = [
@@ -23,15 +23,26 @@ __all__ = [
     "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",
+    "is_master",
     "BaseLauncher",
     "CliLauncher",
     "SingleProcessLauncher",
-    "LogAudio",
     "LogImage",
     "LogLine",
-    "LogVideo",
     "Logger",
     "LoggerImpl",
+    "CallbackLogger",
     "JsonLogger",
     "StateLogger",
     "StdoutLogger",
@@ -46,34 +57,59 @@ __all__ = [
     "collate",
     "collate_non_null",
     "BaseFileDownloader",
+    "CumulativeTimer",
     "DataDownloader",
+    "IntervalTicker",
+    "IterationTimer",
+    "MinGradScaleError",
     "ModelDownloader",
+    "NaNError",
+    "StateTimer",
+    "TrainingFinishedError",
     "check_md5",
     "check_sha256",
+    "cpu_count",
+    "date_str",
+    "diff_configs",
     "get_git_state",
+    "get_random_port",
     "get_state_dict_prefix",
     "get_training_code",
     "save_config",
+    "stage_environment",
+    "to_markdown_table",
     "ColoredFormatter",
     "configure_logging",
     "one_hot",
     "partial_flatten",
     "worker_chunk",
     "TextBlock",
+    "camelcase_to_snakecase",
     "colored",
     "format_datetime",
     "format_timedelta",
+    "highlight_exception_message",
+    "is_interactive_session",
     "outlined",
     "render_text_blocks",
     "show_error",
+    "show_info",
     "show_warning",
+    "snakecase_to_camelcase",
     "uncolored",
     "wrapped",
 ]
 
 __all__ += [
+    "Batch",
     "CollateMode",
+    "EmbeddingKind",
+    "LOG_ERROR_SUMMARY",
+    "LOG_PING",
+    "LOG_STATUS",
+    "Output",
     "Phase",
+    "RawConfigType",
 ]
 
 import os
@@ -95,21 +131,32 @@ NAME_MAP: dict[str, str] = {
     "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",
+    "is_master": "nn.parallel",
     "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",
+    "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",
+    "DataloaderConfig": "task.mixins.data_loader",
     "GPUStatsOptions": "task.mixins.gpu_stats",
     "Script": "task.script",
     "ScriptConfig": "task.script",
@@ -118,27 +165,45 @@ NAME_MAP: dict[str, str] = {
     "collate": "utils.data.collate",
     "collate_non_null": "utils.data.collate",
     "BaseFileDownloader": "utils.experiments",
+    "CumulativeTimer": "utils.experiments",
     "DataDownloader": "utils.experiments",
+    "IntervalTicker": "utils.experiments",
+    "IterationTimer": "utils.experiments",
+    "MinGradScaleError": "utils.experiments",
     "ModelDownloader": "utils.experiments",
+    "NaNError": "utils.experiments",
+    "StateTimer": "utils.experiments",
+    "TrainingFinishedError": "utils.experiments",
     "check_md5": "utils.experiments",
     "check_sha256": "utils.experiments",
+    "cpu_count": "utils.experiments",
+    "date_str": "utils.experiments",
+    "diff_configs": "utils.experiments",
     "get_git_state": "utils.experiments",
+    "get_random_port": "utils.experiments",
     "get_state_dict_prefix": "utils.experiments",
     "get_training_code": "utils.experiments",
     "save_config": "utils.experiments",
+    "stage_environment": "utils.experiments",
+    "to_markdown_table": "utils.experiments",
     "ColoredFormatter": "utils.logging",
     "configure_logging": "utils.logging",
     "one_hot": "utils.numpy",
     "partial_flatten": "utils.numpy",
     "worker_chunk": "utils.numpy",
     "TextBlock": "utils.text",
+    "camelcase_to_snakecase": "utils.text",
     "colored": "utils.text",
     "format_datetime": "utils.text",
     "format_timedelta": "utils.text",
+    "highlight_exception_message": "utils.text",
+    "is_interactive_session": "utils.text",
     "outlined": "utils.text",
     "render_text_blocks": "utils.text",
     "show_error": "utils.text",
+    "show_info": "utils.text",
     "show_warning": "utils.text",
+    "snakecase_to_camelcase": "utils.text",
     "uncolored": "utils.text",
     "wrapped": "utils.text",
 }
@@ -146,8 +211,15 @@ NAME_MAP: dict[str, str] = {
 # 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",
+        "LOG_ERROR_SUMMARY": "utils.logging",
+        "LOG_PING": "utils.logging",
+        "LOG_STATUS": "utils.logging",
+        "Output": "task.mixins.output",
         "Phase": "core.state",
+        "RawConfigType": "task.base",
     },
 )
 
@@ -171,10 +243,27 @@ if IMPORT_ALL or TYPE_CHECKING:
         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.nn.parallel import is_master
+    from xax.task.base import RawConfigType
     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.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
@@ -182,31 +271,56 @@ if IMPORT_ALL or TYPE_CHECKING:
     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,
+        CumulativeTimer,
         DataDownloader,
+        IntervalTicker,
+        IterationTimer,
+        MinGradScaleError,
         ModelDownloader,
+        NaNError,
+        StateTimer,
+        TrainingFinishedError,
         check_md5,
         check_sha256,
+        cpu_count,
+        date_str,
+        diff_configs,
         get_git_state,
+        get_random_port,
         get_state_dict_prefix,
         get_training_code,
         save_config,
+        stage_environment,
+        to_markdown_table,
+    )
+    from xax.utils.logging import (
+        LOG_ERROR_SUMMARY,
+        LOG_PING,
+        LOG_STATUS,
+        ColoredFormatter,
+        configure_logging,
     )
-    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,
+        camelcase_to_snakecase,
         colored,
         format_datetime,
         format_timedelta,
+        highlight_exception_message,
+        is_interactive_session,
         outlined,
         render_text_blocks,
         show_error,
+        show_info,
         show_warning,
+        snakecase_to_camelcase,
         uncolored,
         wrapped,
     )

xax/core/conf.py

@@ -6,7 +6,6 @@ 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
@@ -61,67 +60,44 @@ def is_missing(cfg: Any, key: str) -> bool:  # noqa: ANN401
     return False
 
 
-@dataclass
+@dataclass(kw_only=True)
 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")
+    log_level: str = field(II("oc.env:XAX_LOG_LEVEL,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
+@dataclass(kw_only=True)
 class Triton:
     use_triton_if_available: bool = field(True, help="Use Triton if available")
 
 
-@dataclass
+@dataclass(kw_only=True)
 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
+@dataclass(kw_only=True)
 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
+@dataclass(kw_only=True)
 class SlurmPartition:
     partition: str = field(MISSING, help="The partition name")
     num_nodes: int = field(1, help="The number of nodes to use")
 
 
-@dataclass
+@dataclass(kw_only=True)
 class Slurm:
     launch: dict[str, SlurmPartition] = field({}, help="The available launch configurations")
 
 
-@dataclass
+@dataclass(kw_only=True)
 class UserConfig:
     logging: Logging = field(Logging)
-    device: Device = field(Device)
     triton: Triton = field(Triton)
     experiment: Experiment = field(Experiment)
     directories: Directories = field(Directories)

xax/core/state.py

@@ -2,7 +2,7 @@
 
 import time
 from dataclasses import dataclass
-from typing import Literal, TypedDict, cast, get_args
+from typing import Literal, NotRequired, TypedDict, cast, get_args
 
 from omegaconf import MISSING
 
@@ -18,16 +18,16 @@ def cast_phase(raw_phase: str) -> 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
+    num_steps: NotRequired[int]
+    num_samples: NotRequired[int]
+    num_valid_steps: NotRequired[int]
+    num_valid_samples: NotRequired[int]
+    start_time_s: NotRequired[float]
+    elapsed_time_s: NotRequired[float]
+    raw_phase: NotRequired[str]
 
 
-@dataclass(frozen=True)
+@dataclass
 class State:
     num_steps: int = field(MISSING, help="Number of steps so far")
     num_samples: int = field(MISSING, help="Number of sample so far")
@@ -41,6 +41,10 @@ class State:
     def phase(self) -> Phase:
         return cast_phase(self.raw_phase)
 
+    @phase.setter
+    def phase(self, phase: Phase) -> None:
+        self.raw_phase = phase
+
     @classmethod
     def init_state(cls) -> "State":
         return cls(
@@ -65,17 +69,3 @@ class State:
                 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})