pfnstudio-core 0.7.0__py3-none-any.whl

pfnstudio_core/__init__.py

@@ -0,0 +1,68 @@
+"""PFN Studio core abstractions."""
+
+from . import axes as _axes  # registers built-in axes on import
+from . import blocks as _blocks  # registers built-in blocks on import
+from .axes import (
+    UNKNOWN,
+    Axis,
+    AxisDetector,
+    detect,
+    encode_tag,
+    get_axis,
+    is_unknown,
+    list_axes,
+    register_axis,
+    sample_tag,
+    tag_dim,
+    train_detector,
+)
+from .datasets import DatasetUnavailable, RegistryDatasetLoader
+from .eval import Eval, EvalSpec
+from .model import Model, ModelSpec
+from .prior import Prior, PriorSpec
+from .registry import (
+    get_block,
+    get_eval,
+    get_prior,
+    list_blocks,
+    list_evals,
+    list_priors,
+    register_block,
+    register_eval,
+    register_prior,
+)
+from .run import Run, RunSpec
+
+__version__ = "0.4.0"
+
+__all__ = [
+    "UNKNOWN",
+    "Axis",
+    "AxisDetector",
+    "Eval",
+    "EvalSpec",
+    "Model",
+    "ModelSpec",
+    "Prior",
+    "PriorSpec",
+    "Run",
+    "RunSpec",
+    "detect",
+    "encode_tag",
+    "get_axis",
+    "get_block",
+    "get_eval",
+    "get_prior",
+    "is_unknown",
+    "list_axes",
+    "list_blocks",
+    "list_evals",
+    "list_priors",
+    "register_axis",
+    "register_block",
+    "register_eval",
+    "register_prior",
+    "sample_tag",
+    "tag_dim",
+    "train_detector",
+]

pfnstudio_core/axes/__init__.py

@@ -0,0 +1,34 @@
+"""Axes package — public surface for promptable-prior axes.
+
+Importing this module also imports `builtins`, which eagerly
+registers the built-in axes (monotonicity, …) so they're available
+without an extra import in user code."""
+
+from . import builtins as _builtins  # noqa: F401  (side-effect: registers built-ins)
+from .base import (
+    UNKNOWN,
+    Axis,
+    AxisKind,
+    get_axis,
+    is_unknown,
+    list_axes,
+    register_axis,
+)
+from .detector import AxisDetector, detect, train_detector
+from .encoding import encode_tag, sample_tag, tag_dim
+
+__all__ = [
+    "UNKNOWN",
+    "Axis",
+    "AxisDetector",
+    "AxisKind",
+    "detect",
+    "encode_tag",
+    "get_axis",
+    "is_unknown",
+    "list_axes",
+    "register_axis",
+    "sample_tag",
+    "tag_dim",
+    "train_detector",
+]

pfnstudio_core/axes/base.py

@@ -0,0 +1,111 @@
+"""Axis abstraction for promptable priors.
+
+An Axis declares a *steerable property* of a prior — a knob the trained
+brain learns to honor at inference time. Examples: monotonicity (edge
+signs in an SCM), lag_scale (delay distribution), feedback_allowed
+(DAG vs cyclic), sparsity (mean parents per node).
+
+The axis declares the contract (name, kind, value space, default).
+Each Prior that opts into an axis implements its own application
+logic — the same axis can mean different things to different priors
+(e.g. monotonicity on an SCM constrains edge signs; on a regression
+prior it constrains weight signs). The base class doesn't try to
+unify those interpretations.
+
+Back-compat invariants this module enforces:
+- The reserved sentinel ``UNKNOWN`` means "skip this axis at sample
+  time". A prior given ``tag={axis: UNKNOWN}`` for every axis is
+  bit-identical to its pre-axis behavior. This is what guarantees
+  adding axes can't regress existing benchmarks.
+- A prior with no declared axes ignores any ``tag=...`` argument
+  passed to ``sample()`` — old call sites keep working unchanged.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+AxisKind = Literal["categorical", "range", "boolean"]
+
+# Reserved sentinel — when a tag has this for an axis, the prior must
+# skip the axis hook and behave as if the axis weren't declared.
+# Stored as a string so it round-trips through JSON / YAML cleanly.
+UNKNOWN = "__unknown__"
+
+
+@dataclass(frozen=True)
+class Axis:
+    name: str
+    kind: AxisKind
+    # For categorical: list of label strings. For range: [min, max].
+    # For boolean: ignored.
+    values: tuple[Any, ...] = field(default_factory=tuple)
+    description: str = ""
+    # Fraction of training samples that get UNKNOWN for this axis.
+    # Default 0.3 — preserves substantial unconditional mass so the
+    # brain doesn't lose its pre-axis baseline.
+    unknown_mass: float = 0.3
+
+    def __post_init__(self) -> None:
+        if not 0.0 <= self.unknown_mass <= 1.0:
+            raise ValueError(
+                f"axis {self.name!r}: unknown_mass must be in [0, 1], got {self.unknown_mass}"
+            )
+        if self.kind == "categorical" and not self.values:
+            raise ValueError(f"axis {self.name!r}: categorical axes need at least one value")
+        if self.kind == "range" and len(self.values) != 2:
+            raise ValueError(
+                f"axis {self.name!r}: range axes need exactly 2 values (min, max), got {len(self.values)}"
+            )
+
+    def sample_value(self, rng: Any) -> Any:
+        """Sample one value for this axis, honoring the unknown_mass invariant."""
+        if rng.random() < self.unknown_mass:
+            return UNKNOWN
+        if self.kind == "categorical":
+            return rng.choice(list(self.values))
+        if self.kind == "boolean":
+            return bool(rng.integers(0, 2))
+        if self.kind == "range":
+            lo, hi = float(self.values[0]), float(self.values[1])
+            return float(rng.uniform(lo, hi))
+        raise ValueError(f"unknown axis kind: {self.kind!r}")
+
+
+# Module-level registry. Resolves axis names → Axis instances at
+# loader time. Built-in axes register themselves on import; custom
+# axes register via @register_axis.
+_AXES: dict[str, Axis] = {}
+
+
+def register_axis(axis: Axis) -> Axis:
+    """Register an Axis. Idempotent: re-registering the *same* axis is
+    a no-op (useful when modules get imported twice); re-registering
+    a *different* axis under the same name raises."""
+    existing = _AXES.get(axis.name)
+    if existing is not None and existing != axis:
+        raise ValueError(
+            f"axis {axis.name!r} already registered with a different definition"
+        )
+    _AXES[axis.name] = axis
+    return axis
+
+
+def get_axis(name: str) -> Axis:
+    """Look up a registered axis by name. Raises KeyError if not registered."""
+    if name not in _AXES:
+        raise KeyError(
+            f"axis {name!r} not registered. Known axes: {sorted(_AXES)}"
+        )
+    return _AXES[name]
+
+
+def list_axes() -> list[str]:
+    return sorted(_AXES)
+
+
+def is_unknown(value: Any) -> bool:
+    """Whether a tag value is the UNKNOWN sentinel. Use this in axis
+    application code rather than comparing to the string literal."""
+    return value == UNKNOWN

pfnstudio_core/axes/builtins/__init__.py

@@ -0,0 +1,5 @@
+"""Built-in axes. Importing this module registers them with the global
+axis registry. The package ``__init__`` imports this so axes are
+available the moment ``pfnstudio_core`` is imported."""
+
+from . import monotonicity  # noqa: F401  (registers on import)

pfnstudio_core/axes/builtins/monotonicity.py

@@ -0,0 +1,25 @@
+"""Built-in `monotonicity` axis.
+
+Declares the contract; per-prior application lives in each Prior
+subclass (an SCM's monotonicity constrains edge signs, a regression
+prior's monotonicity constrains weight signs, etc.).
+"""
+
+from __future__ import annotations
+
+from ..base import Axis, register_axis
+
+monotonicity = register_axis(
+    Axis(
+        name="monotonicity",
+        kind="categorical",
+        values=("positive", "negative", "mixed"),
+        description=(
+            "Sign discipline of the prior's causal relationships. "
+            "positive = X up implies Y up everywhere; negative = "
+            "X up implies Y down; mixed = signs may vary. Unknown "
+            "(default) leaves the prior unconstrained."
+        ),
+        unknown_mass=0.3,
+    )
+)

pfnstudio_core/axes/detector.py

@@ -0,0 +1,260 @@
+"""Auto-detector — supervised model that reads context and proposes
+axis values.
+
+The pitch in one line: instead of asking the user to fill out an empty
+form of axis chips, the detector reads their data and pre-fills the
+chips with its best guess (and a confidence). The user just confirms,
+overrides, or adds the axes the detector can't see.
+
+Training:
+- Each batch, the prior samples a real (non-UNKNOWN) tag value and
+  generates data from it. The tag value is the supervised label.
+- The detector reads the data (no tag) and predicts a softmax over
+  axis values.
+- Cross-entropy loss.
+
+Inference:
+- Same encoder forward, output softmax probabilities.
+- Highest-probability value becomes the proposed chip; that probability
+  becomes the displayed confidence.
+
+Limitations in v1:
+- Only categorical axes are detected. Range axes need a regression
+  head and a different loss; boolean axes need a 2-class head.
+  Both are straightforward additions but defer until they're actually
+  shipped on a prior.
+- The detector is small (32-d, 2-layer) — same shape as the test
+  brain. Scale up when a real prior with many axes ships.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from .base import UNKNOWN, Axis
+
+
+class AxisDetector:
+    """A small encoder + per-axis classification head. Holds one
+    head per categorical axis; non-categorical axes are silently
+    skipped (an empty detector is legal — predicts nothing)."""
+
+    def __init__(
+        self,
+        axes: list[Axis],
+        d_model: int = 32,
+        n_heads: int = 4,
+        n_layers: int = 2,
+        dropout: float = 0.0,
+    ):
+        try:
+            import torch.nn as nn
+        except ImportError as e:
+            raise ImportError(
+                "AxisDetector requires torch. "
+                "Install with: pip install pfnstudio-core[torch]"
+            ) from e
+
+        # Only categorical axes get heads in v1. Skipped axes are
+        # tracked so callers can warn the user "we can't propose for
+        # axis X yet."
+        self._scored_axes: list[Axis] = [
+            a for a in axes if a.kind == "categorical" and len(a.values) >= 2
+        ]
+        self._skipped_axes: list[Axis] = [
+            a for a in axes if a not in self._scored_axes
+        ]
+        self.d_model = d_model
+
+        # Lazy embedder — input feature dim is inferred on first
+        # forward. Same pattern the TabularEmbedder uses.
+        self.embedder = nn.LazyLinear(d_model)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=n_heads,
+            dim_feedforward=d_model * 4,
+            dropout=dropout,
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
+        # One head per scored axis. Output dim = number of values
+        # (no UNKNOWN class — the detector proposes a real value;
+        # low confidence across all classes is how it expresses
+        # "I'm not sure", not a separate class).
+        self.heads = nn.ModuleDict(
+            {a.name: nn.Linear(d_model, len(a.values)) for a in self._scored_axes}
+        )
+
+    @property
+    def scored_axes(self) -> list[Axis]:
+        """Axes the detector has heads for (categorical, ≥2 values)."""
+        return self._scored_axes
+
+    @property
+    def skipped_axes(self) -> list[Axis]:
+        """Axes the detector can't predict (boolean, range, single-value)."""
+        return self._skipped_axes
+
+    def parameters(self):
+        # nn.ModuleDict + nn.Module attributes — yield everything.
+        yield from self.embedder.parameters()
+        yield from self.encoder.parameters()
+        yield from self.heads.parameters()
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "embedder": self.embedder.state_dict(),
+            "encoder": self.encoder.state_dict(),
+            "heads": self.heads.state_dict(),
+        }
+
+    def __call__(self, x: Any) -> dict[str, Any]:
+        """Forward pass. Returns axis_name → logits tensor of shape
+        (B, num_values_for_axis)."""
+        emb = self.embedder(x)  # (B, N, d_model)
+        encoded = self.encoder(emb)  # (B, N, d_model)
+        # Mean-pool the context tokens to one vector per batch element.
+        # (Could swap for an attention pool later; mean is fine for
+        # the monotonicity demo and avoids extra parameters.)
+        pooled = encoded.mean(dim=1)  # (B, d_model)
+        return {name: head(pooled) for name, head in self.heads.items()}
+
+
+def _value_to_index(axis: Axis, value: Any) -> int:
+    """Index of `value` in axis.values. Raises if value is UNKNOWN
+    or not a known value — UNKNOWN tags shouldn't be used as training
+    labels because the detector is supposed to learn to assign a
+    concrete value."""
+    if value == UNKNOWN:
+        raise ValueError(
+            f"axis {axis.name!r}: cannot train detector on UNKNOWN label"
+        )
+    return list(axis.values).index(value)
+
+
+def train_detector(
+    *,
+    detector: AxisDetector,
+    prior: Any,
+    steps: int = 500,
+    batch_size: int = 16,
+    lr: float = 1e-3,
+    seed: int = 0,
+    on_step: Any = None,
+) -> dict[str, Any]:
+    """Train the detector on labeled batches sampled from the prior.
+
+    Each batch picks a non-UNKNOWN value for each scored axis, samples
+    data with that tag, and trains the detector to predict the tag
+    from the data. Cross-entropy loss, summed across axes.
+
+    Returns ``{status, final_loss, mean_acc}`` where ``mean_acc`` is
+    the accuracy of the detector on the *last* batch (cheap sanity
+    check that training did something)."""
+    try:
+        import torch
+        import torch.nn.functional as F
+    except ImportError:
+        return {"status": "skipped", "reason": "torch not installed"}
+
+    if not detector.scored_axes:
+        return {"status": "skipped", "reason": "detector has no scored axes"}
+
+    rng = np.random.default_rng(seed)
+    optim = torch.optim.AdamW(list(detector.parameters()), lr=lr)
+
+    last_acc: dict[str, float] = {}
+    final_loss = 0.0
+    for step in range(steps):
+        # Sample one tag value per axis (random, non-UNKNOWN — that's
+        # the supervised label the detector learns to recover).
+        tag = {
+            a.name: rng.choice(list(a.values)) for a in detector.scored_axes
+        }
+        batch = prior.sample_batch(
+            batch_size=batch_size, seed=seed + step * batch_size, tag=tag
+        )
+        X = torch.stack([torch.from_numpy(b["X"]).float() for b in batch])
+
+        logits = detector(X)  # dict axis → (B, num_values)
+        loss = torch.zeros(1)
+        for axis in detector.scored_axes:
+            label_idx = _value_to_index(axis, tag[axis.name])
+            labels = torch.full((X.shape[0],), label_idx, dtype=torch.long)
+            loss = loss + F.cross_entropy(logits[axis.name], labels)
+
+        optim.zero_grad()
+        loss.backward()
+        optim.step()
+
+        final_loss = float(loss.item())
+        if on_step is not None:
+            on_step(step, final_loss)
+
+        # Final-batch accuracy (per axis) — cheap and informative.
+        if step == steps - 1:
+            with torch.no_grad():
+                for axis in detector.scored_axes:
+                    preds = logits[axis.name].argmax(dim=-1)
+                    label_idx = _value_to_index(axis, tag[axis.name])
+                    last_acc[axis.name] = float((preds == label_idx).float().mean().item())
+
+    return {
+        "status": "ok",
+        "steps": steps,
+        "final_loss": final_loss,
+        "last_batch_accuracy": last_acc,
+    }
+
+
+def detect(
+    *,
+    detector: AxisDetector,
+    context: Any,
+) -> dict[str, dict[str, Any]]:
+    """Run the detector on a context tensor and return a tag-shaped
+    proposal with per-axis confidence.
+
+    Returns ``{axis_name: {value, confidence, probs}}`` where:
+      - ``value`` is the predicted axis value (argmax)
+      - ``confidence`` is the softmax probability of that value
+      - ``probs`` is the full {value: probability} dict, for callers
+        that want to render uncertainty (e.g. show "60% positive,
+        35% mixed, 5% negative" instead of just "positive 60%").
+    """
+    try:
+        import torch
+        import torch.nn.functional as F
+    except ImportError:
+        return {}
+
+    if not detector.scored_axes:
+        return {}
+
+    if not torch.is_tensor(context):
+        context = torch.from_numpy(np.asarray(context)).float()
+    if context.dim() == 2:
+        context = context.unsqueeze(0)  # add batch dim
+
+    with torch.no_grad():
+        logits = detector(context)
+
+    out: dict[str, dict[str, Any]] = {}
+    for axis in detector.scored_axes:
+        probs = F.softmax(logits[axis.name], dim=-1)
+        # Mean across batch in case the caller passed multiple contexts —
+        # the detector returns one proposal per axis, not per-sample.
+        mean_probs = probs.mean(dim=0)
+        best_idx = int(mean_probs.argmax().item())
+        out[axis.name] = {
+            "value": axis.values[best_idx],
+            "confidence": float(mean_probs[best_idx].item()),
+            "probs": {
+                axis.values[i]: float(mean_probs[i].item())
+                for i in range(len(axis.values))
+            },
+        }
+    return out

pfnstudio_core/axes/encoding.py

@@ -0,0 +1,119 @@
+"""Tag encoding — fixed-length vector representation of an axis tag.
+
+The encoding rules:
+
+- **Categorical axes** → one-hot over the declared values, plus one
+  extra slot for the UNKNOWN sentinel (always last). So a categorical
+  axis with K values contributes K+1 dimensions.
+- **Boolean axes** → 2 dims (false, true) plus 1 for UNKNOWN. Same
+  shape as a 2-value categorical.
+- **Range axes** → 2 dims: ``[is_unknown_flag, normalised_value]``.
+  When tagged with a real value v in [min, max], normalised_value is
+  ``(v - min) / (max - min)`` clipped to [0, 1] and is_unknown_flag
+  is 0. When UNKNOWN, normalised_value is 0 and is_unknown_flag is 1.
+
+The encoding is deterministic and stateless — the same axes list and
+the same tag always produce the same vector. ``tag_dim(axes)`` returns
+the total length so model constructors can size their tag embedder
+once at init time without poking at sample tags.
+
+Back-compat invariant: ``encode_tag({}, axes)`` and ``encode_tag(None,
+axes)`` are byte-identical to ``encode_tag({a: UNKNOWN for a in axes},
+axes)``. Used by callers who need a "no constraints" tag vector
+without enumerating all axes.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from .base import UNKNOWN, Axis, is_unknown
+
+
+def _axis_dim(axis: Axis) -> int:
+    if axis.kind == "categorical":
+        return len(axis.values) + 1  # values + UNKNOWN slot
+    if axis.kind == "boolean":
+        return 3  # false, true, UNKNOWN
+    if axis.kind == "range":
+        return 2  # is_unknown_flag, normalised_value
+    raise ValueError(f"unknown axis kind: {axis.kind!r}")
+
+
+def tag_dim(axes: list[Axis]) -> int:
+    """Total length of the tag vector produced by ``encode_tag`` for
+    the given axes (in the same order)."""
+    return sum(_axis_dim(a) for a in axes)
+
+
+def _encode_categorical(axis: Axis, value: Any) -> np.ndarray:
+    """One-hot over values + final UNKNOWN slot."""
+    dim = _axis_dim(axis)
+    out = np.zeros(dim, dtype=np.float32)
+    if value is None or is_unknown(value):
+        out[-1] = 1.0
+        return out
+    if value not in axis.values:
+        raise ValueError(
+            f"axis {axis.name!r}: value {value!r} not in {axis.values}"
+        )
+    idx = axis.values.index(value)
+    out[idx] = 1.0
+    return out
+
+
+def _encode_boolean(axis: Axis, value: Any) -> np.ndarray:
+    out = np.zeros(3, dtype=np.float32)
+    if value is None or is_unknown(value):
+        out[2] = 1.0
+        return out
+    out[int(bool(value))] = 1.0
+    return out
+
+
+def _encode_range(axis: Axis, value: Any) -> np.ndarray:
+    out = np.zeros(2, dtype=np.float32)
+    if value is None or is_unknown(value):
+        out[0] = 1.0  # is_unknown_flag
+        return out
+    lo, hi = float(axis.values[0]), float(axis.values[1])
+    span = hi - lo
+    if span == 0:
+        out[1] = 0.0
+    else:
+        normalised = (float(value) - lo) / span
+        out[1] = float(np.clip(normalised, 0.0, 1.0))
+    return out
+
+
+def encode_tag(tag: dict[str, Any] | None, axes: list[Axis]) -> np.ndarray:
+    """Encode a tag dict into a fixed-length float32 vector for the given axes.
+
+    Missing axis entries are treated as UNKNOWN — same as explicitly
+    passing UNKNOWN. The encoded vector's length matches ``tag_dim(axes)``.
+    """
+    tag = tag or {}
+    chunks: list[np.ndarray] = []
+    for axis in axes:
+        value = tag.get(axis.name, UNKNOWN)
+        if axis.kind == "categorical":
+            chunks.append(_encode_categorical(axis, value))
+        elif axis.kind == "boolean":
+            chunks.append(_encode_boolean(axis, value))
+        elif axis.kind == "range":
+            chunks.append(_encode_range(axis, value))
+        else:
+            raise ValueError(f"unknown axis kind: {axis.kind!r}")
+    if not chunks:
+        return np.zeros(0, dtype=np.float32)
+    return np.concatenate(chunks, axis=0)
+
+
+def sample_tag(axes: list[Axis], rng: np.random.Generator) -> dict[str, Any]:
+    """Sample one tag dict for all axes, honoring each axis's
+    ``unknown_mass``. Returns axis_name → sampled value (or UNKNOWN).
+
+    Used by the training loop to draw a fresh tag per batch."""
+    return {axis.name: axis.sample_value(rng) for axis in axes}