dewml 0.1.0__py3-none-any.whl

dew/__init__.py

@@ -0,0 +1,143 @@
+"""Dew: one objective, one trainer.
+
+Each name exported here is imported from its own module when you first
+access it, not at `import dew`. So `import dew.training` loads only the
+training layer, with no modality, encoder or tracker backend. Importing
+`dew` opens no JAX backend and loads no optional dependency; encoders,
+decoders and datasets load what they need when they are built.
+
+`import dew` does set two XLA flags before the backend opens.
+`--xla_allow_excess_precision=false` makes XLA round values declared in a
+narrow dtype such as bf16 where the program rounds them. When JAX's CUDA
+plugin is installed, `--xla_gpu_enable_allocator_spatial_partitioning=false`
+keeps a preallocated GPU pool in one piece for a step's temporaries. If you
+set either flag yourself in XLA_FLAGS, your value is kept. If the JAX
+backend has already opened, the flags cannot take effect, and Dew logs a
+warning.
+"""
+
+from collections.abc import Callable
+from importlib import import_module
+from typing import TYPE_CHECKING
+
+from dew.logging import configure as _configure_logging
+from dew.telemetry.devices import (
+    keep_roundings as _keep_roundings,
+    unpartition_gpu_pool as _unpartition_gpu_pool,
+)
+
+_configure_logging()
+_keep_roundings()
+_unpartition_gpu_pool()
+
+if TYPE_CHECKING:  # the surface above, with its types, for checkers and editors
+    from dew.artifacts import ImageGrid, Representations, TextSamples, TokenScores, VideoGrid
+    from dew.data import Dataset
+    from dew.diffusion import Process
+    from dew.eval import Mean
+    from dew.inference import pipeline
+    from dew.inputs import Condition, Field, InputSpec
+    from dew.objectives import Objective
+    from dew.objectives.base import Aux, EMASpec, Step
+    from dew.objectives.supervised import Supervised
+    from dew.sampling import CFG, sample
+    from dew.telemetry.profile import Profiler
+    from dew.training import (
+        Best,
+        Checkpoints,
+        EvalSuite,
+        Evaluation,
+        Keep,
+        Layout,
+        LocalTracker,
+        MeshSpec,
+        MLflowTracker,
+        Plateau,
+        ProfileWindow,
+        TensorBoardTracker,
+        Tracker,
+        Trackers,
+        Trainer,
+        TrainState,
+        WandbTracker,
+    )
+
+__version__ = "0.1.0"
+
+_EXPORTS = {
+    "Best": "dew.training", "Keep": "dew.training", "Plateau": "dew.training", "EvalSuite": "dew.training",
+    "Trainer": "dew.training", "TrainState": "dew.training", "Step": "dew.training",
+    "Aux": "dew.training", "EMASpec": "dew.training", "MeshSpec": "dew.training",
+    "Layout": "dew.training", "Checkpoints": "dew.training", "Tracker": "dew.training",
+    "WandbTracker": "dew.training", "LocalTracker": "dew.training", "Trackers": "dew.training",
+    "MLflowTracker": "dew.training", "TensorBoardTracker": "dew.training",
+    "Evaluation": "dew.training",
+    "ProfileWindow": "dew.training",
+    "Objective": "dew.objectives", "Supervised": "dew.objectives.supervised",
+    "Dataset": "dew.data",
+    "Process": "dew.diffusion",
+    "Mean": "dew.eval",
+    "InputSpec": "dew.inputs", "Field": "dew.inputs", "Condition": "dew.inputs",
+    "sample": "dew.sampling", "CFG": "dew.sampling",
+    "pipeline": "dew.inference",
+    "Profiler": "dew.telemetry.profile",
+    "ImageGrid": "dew.artifacts", "VideoGrid": "dew.artifacts",
+    "TextSamples": "dew.artifacts", "Representations": "dew.artifacts",
+    "TokenScores": "dew.artifacts",
+}
+
+
+def __getattr__(name: str) -> type | Callable:
+    module = _EXPORTS.get(name)
+    if module is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    return getattr(import_module(module), name)
+
+
+def __dir__() -> list[str]:
+    return list(__all__)
+
+
+# Written out, not derived from _EXPORTS, so a type checker, an editor and
+# `from dew import *` can all read the public surface without running the
+# lazy lookup above.
+__all__ = [
+    "CFG",
+    "Aux",
+    "Best",
+    "Checkpoints",
+    "Condition",
+    "Dataset",
+    "EMASpec",
+    "EvalSuite",
+    "Evaluation",
+    "Field",
+    "ImageGrid",
+    "InputSpec",
+    "Keep",
+    "Layout",
+    "LocalTracker",
+    "MLflowTracker",
+    "Mean",
+    "MeshSpec",
+    "Objective",
+    "Plateau",
+    "Process",
+    "ProfileWindow",
+    "Profiler",
+    "Representations",
+    "Step",
+    "Supervised",
+    "TensorBoardTracker",
+    "TextSamples",
+    "TokenScores",
+    "Tracker",
+    "Trackers",
+    "TrainState",
+    "Trainer",
+    "VideoGrid",
+    "WandbTracker",
+    "__version__",
+    "pipeline",
+    "sample",
+]

dew/_model_types.py

@@ -0,0 +1,9 @@
+"""Model-type spellings shared by the Qwen decoder, wrapper and vision maps.
+
+Vision translation reads these without importing the interop package, whose
+public entry point imports the vision classes itself.
+"""
+
+QWEN35_TYPES = ('qwen3_5', 'qwen3_5_moe')
+QWEN35_TEXT_TYPES = tuple(f'{name}_text' for name in QWEN35_TYPES)
+_QWEN35_VISION_TYPES = (*QWEN35_TYPES, *(f'{name}_vision' for name in QWEN35_TYPES))

dew/artifacts.py

@@ -0,0 +1,117 @@
+"""Typed values that an objective's evaluation produces.
+
+An objective returns these from its scoring and preview hooks. A metric
+reads the scoring artifact of the type it expects, and a tracker renders
+each preview according to its type. The array fields are pytree leaves, so
+they can pass through jit, while the optional captions and decoded text stay
+on the host as static metadata.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import jax
+import numpy as np
+from flax import struct
+from jax.typing import ArrayLike
+from numpy.typing import NDArray
+
+if TYPE_CHECKING:
+    from _typeshed import DataclassInstance
+
+
+@struct.dataclass
+class ImageGrid:
+    """Images in [-1, 1], `[N, H, W, C]`, with the text each was conditioned on
+    where there was any."""
+    images: jax.Array
+    captions: tuple[str, ...] = struct.field(pytree_node=False, default=())
+
+
+@struct.dataclass
+class VideoGrid:
+    """Clips in [-1, 1], `[N, T, H, W, C]`, with the text each was conditioned on
+    where there was any."""
+    videos: jax.Array
+    captions: tuple[str, ...] = struct.field(pytree_node=False, default=())
+
+
+@struct.dataclass
+class TextSamples:
+    """Generated token rows, with optional decoded preview text and prompt."""
+    tokens: jax.Array | np.ndarray
+    prompt: str = struct.field(pytree_node=False, default="")
+    texts: tuple[str, ...] = struct.field(pytree_node=False, default=())
+
+
+@struct.dataclass
+class Representations:
+    """Encoder outputs `[N, D]` and the labels of the records they came from,
+    for a probe to score."""
+    features: jax.Array
+    labels: jax.Array
+
+
+@struct.dataclass
+class TokenScores:
+    """Teacher-forced per-token losses `[N, L]` and the weight of each target.
+
+    A weight is 1 where the target counts and 0 where it is padding or a
+    document's first token. A perplexity is exp of the weighted mean loss
+    over a whole pass, so a batch with no counted target adds nothing to it.
+    """
+    losses: jax.Array
+    weights: jax.Array
+    correct: jax.Array
+    """Per-token top-1 correctness, from the same logits that produced the losses."""
+
+
+@struct.dataclass
+class Decisions:
+    """A decision model's probabilities `[N, Q, K]` over the options of each
+    row's questions, the real options `[N, Q, K]`, each question's right option
+    `[N, Q]`, which questions `[N, Q]` ask a score, whose options are ordered
+    levels, and which `[N, Q]` have a known answer to score."""
+    probabilities: jax.Array
+    options: jax.Array
+    labels: jax.Array
+    ordinal: jax.Array
+    scored: jax.Array
+
+
+type Artifact = DataclassInstance
+"""What an objective's scoring and preview hooks return: a dataclass whose
+per-row fields lead with the batch's rows, which a validation pass cuts to
+the real ones. Dew's own are the classes above. A package's objective may
+score into a dataclass of its own, which its metrics read by type
+(`Metric.reads`); a metric picks exactly one scoring artifact, and previews
+never satisfy metrics. A tracker shows a preview of Dew's types, so a
+package shows its own as one of them, a spike raster as an `ImageGrid`."""
+
+type Artifacts = Artifact | tuple[Artifact, ...]
+"""One artifact, or several."""
+
+
+def uint8_pixels(images: ArrayLike) -> NDArray[np.uint8]:
+    """Convert [-1, 1] pixels, as `ImageGrid` and `VideoGrid` hold them, to uint8 in [0, 255].
+
+    Each value maps to its nearest level, computed as `(x + 1) * 127.5` in
+    float32 and rounded half to even (`np.rint`). The result is clipped,
+    because a sample can leave the range. Metrics score these bytes and
+    trackers preview them, so both see the same image.
+    """
+    levels = np.rint((np.asarray(images, np.float32) + 1.0) * 127.5)
+    return np.clip(levels, 0, 255).astype(np.uint8)
+
+
+__all__ = [
+    "Artifact",
+    "Decisions",
+    "ImageGrid",
+    "Representations",
+    "TextSamples",
+    "TokenScores",
+    "VideoGrid",
+    "uint8_pixels",
+]

dew/cache.py

@@ -0,0 +1,65 @@
+"""Where Dew keeps what it caches on disk, and JAX's persistent compilation cache.
+
+Kept apart from the telemetry and the loaders that read it, so the
+interop, config and inference modules take their cache paths without
+importing the FLOP accounting.
+"""
+
+import os
+import sys
+
+import jax
+
+
+def dew_cache_dir() -> str:
+    """Dew's cache directory: `$XDG_CACHE_HOME/dew`, else ~/.cache/dew."""
+    return os.path.expanduser(
+        os.path.join(os.environ.get("XDG_CACHE_HOME") or os.path.join("~", ".cache"), "dew")
+    )
+
+
+def default_compilation_cache_dir() -> str:
+    """Where compiled executables go unless a run names somewhere else.
+
+    The directory JAX is configured with (`jax_compilation_cache_dir`, which
+    JAX_COMPILATION_CACHE_DIR sets) when there is one, so a machine keeps one
+    cache for every entry point. Otherwise Python minors have separate
+    defaults: jax 0.11.2 compresses with Python 3.14's stdlib zstd but names
+    the codec "zlib" in the key, which says "zstandard" only for the
+    zstandard package (`jax._src.compilation_cache.get_cache_key`), so an
+    older interpreter sharing the directory would read those bytes with the
+    wrong codec. Explicit paths passed to enable_compilation_cache remain
+    unchanged.
+    """
+    if jax.config.jax_compilation_cache_dir:
+        return jax.config.jax_compilation_cache_dir
+    return os.path.join(dew_cache_dir(), 'xla', f"python{sys.version_info.major}.{sys.version_info.minor}")
+
+
+def enable_compilation_cache(path: str):
+    """Persist compiled executables so restarts skip XLA compilation.
+
+    The dominant cost of a restart-heavy TPU workflow, where every run otherwise
+    recompiles the same step function from scratch.
+    """
+    os.makedirs(path, exist_ok=True)
+    jax.config.update('jax_compilation_cache_dir', path)
+    # Defaults skip small/fast compilations; a training step is neither, and
+    # caching everything keeps startup predictable.
+    jax.config.update('jax_persistent_cache_min_entry_size_bytes', -1)
+    jax.config.update('jax_persistent_cache_min_compile_time_secs', 0.0)
+
+
+def persist_compilations() -> None:
+    """Point XLA at the on-disk executable cache, unless a directory is set.
+
+    A loaded task compiles for seconds the first time its shapes are seen
+    (a minute for a text-to-image sample on an A100), and a serving process
+    restarts. Training turns the same cache on in `prepare_process`; a task
+    turns it on where it is loaded, `dew.pipeline` or a saved run's record
+    (`dew.inference.tasks.run_record`). Reading the setting is what makes it
+    idempotent and what leaves a trainer's own directory, or a caller's, alone.
+    """
+    if jax.config.jax_compilation_cache_dir:
+        return
+    enable_compilation_cache(default_compilation_cache_dir())