causilo 1.0.0__py3-none-any.whl

causilo/__init__.py

@@ -0,0 +1,7 @@
+"""Tabular prediction with fixed pretrained context models."""
+
+__version__ = "1.0.0"
+
+from .estimators import CausiloClassifier, CausiloRegressor
+
+__all__ = ["CausiloClassifier", "CausiloRegressor"]

causilo/checkpoints.py

@@ -0,0 +1,54 @@
+"""Load pinned safetensors weights and verify their task/configuration pairing."""
+
+import hashlib
+import json
+from pathlib import Path
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors import safe_open
+
+from .model import Model, ModelConfig
+
+REPOSITORY = "nums-ai/causilo"
+RELEASE_COMMIT = "94f2bd91db0737d4da59f347910662905ecb5a09"
+
+
+def load_checkpoint(directory: Path) -> Model:
+    """Validate the format and config hash, then materialize an evaluation model."""
+    encoded = (directory / "config.json").read_bytes()
+    record = json.loads(encoded)
+    expected = {"architecture": "causilo-v1.0", "format_version": 1, "config_version": 1}
+    if any(record.get(key) != value for key, value in expected.items()):
+        raise ValueError("Unsupported model configuration")
+    config = ModelConfig(**record["model"])
+    with safe_open(directory / "model.safetensors", framework="pt") as artifact:
+        metadata = artifact.metadata()
+        required = {
+            "architecture": "causilo-v1.0",
+            "format_version": "1",
+            "task": config.task,
+            "config_sha256": hashlib.sha256(encoded).hexdigest(),
+        }
+        if metadata != required:
+            raise ValueError("Weights do not match their task and configuration")
+        tensors = {key: artifact.get_tensor(key) for key in artifact.keys()}
+    # Construct shapes without allocating a second full copy of the weights.
+    with torch.device("meta"):
+        model = Model(config)
+    model.load_state_dict(tensors, strict=True, assign=True)
+    return model.eval()
+
+
+def load_pretrained_model(task: str) -> Model:
+    """Fetch the task's files at the pinned revision, reusing the Hub cache when available."""
+    folder = "classifier" if task == "classification" else "regressor"
+    root = snapshot_download(
+        REPOSITORY,
+        revision=RELEASE_COMMIT,
+        allow_patterns=[f"{folder}/config.json", f"{folder}/model.safetensors"],
+    )
+    model = load_checkpoint(Path(root) / folder)
+    if model.config.task != task:
+        raise ValueError("The official checkpoint has the wrong prediction task")
+    return model

causilo/data/__init__.py

@@ -0,0 +1 @@
+"""Fitted tabular representations and deterministic ensemble views."""

causilo/data/dataset.py

@@ -0,0 +1,108 @@
+"""Fitted table transforms and ensemble metadata shared by both prediction paths."""
+
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+from sklearn.preprocessing import LabelEncoder, StandardScaler
+from sklearn.utils.multiclass import type_of_target
+
+from .encoding import FeatureEncoder
+from .ensemble import EnsembleMember, make_ensemble_members
+from .normalization import Normalizer
+
+
+@dataclass
+class PreparedDataset:
+    """Training context in encoded feature space, before member permutations.
+
+    ``features`` keeps the encoded training table even when transformed tables
+    are not retained. ``members`` is in construction order; execution and saved
+    K/V caches use ``members_by_normalization()`` order instead. Field names are
+    part of the fitted-state serialization format.
+    """
+
+    encoder: FeatureEncoder
+    features: np.ndarray
+    targets: np.ndarray
+    target_encoder: LabelEncoder | StandardScaler
+    members: tuple[EnsembleMember, ...]
+    normalizers: dict[str, Normalizer]
+    normalized_cache: dict[str, np.ndarray]
+
+    @classmethod
+    def prepare(
+        cls,
+        table,
+        targets,
+        *,
+        task: str,
+        n_estimators: int,
+        retain_preprocessing: bool,
+        max_classes: int,
+        random_state: int,
+    ) -> "PreparedDataset":
+        """Fit transforms on training data only; ``max_classes`` applies to classification."""
+        labels = np.asarray(targets)
+        if labels.ndim == 2 and labels.shape[1] == 1:
+            labels = labels[:, 0]
+        if labels.ndim != 1 or len(labels) != len(table):
+            raise ValueError("Targets must contain one value per training row")
+        if pd.isna(labels).any():
+            raise ValueError("Targets cannot contain missing values")
+        classes = 0
+        if task == "classification":
+            if type_of_target(labels) not in {"binary", "multiclass"}:
+                raise ValueError("Classification targets must be discrete labels")
+            encoder = LabelEncoder().fit(labels)
+            classes = len(encoder.classes_)
+            if not 1 <= classes <= max_classes:
+                raise ValueError(f"Found {classes} classes; this checkpoint supports at most {max_classes}")
+            encoded_targets = encoder.transform(labels)
+        else:
+            # Center in float64 before reducing precision for model execution.
+            labels = labels.astype(np.float64)
+            if not np.isfinite(labels).all():
+                raise ValueError("Regression targets must be finite")
+            encoder = StandardScaler().fit(labels[:, None])
+            encoded_targets = encoder.transform(labels[:, None])[:, 0].astype(np.float32)
+        schema, training = FeatureEncoder.fit(table)
+        members = make_ensemble_members(training.shape[1], classes, n_estimators, random_state)
+        methods = dict.fromkeys(member.normalization for member in members)
+        transforms = {method: Normalizer.fit(training, method) for method in methods}
+        retained = (
+            {method: transform.transform(training) for method, transform in transforms.items()}
+            if retain_preprocessing
+            else {}
+        )
+        return cls(
+            encoder=schema,
+            features=training,
+            targets=encoded_targets,
+            target_encoder=encoder,
+            members=members,
+            normalizers=transforms,
+            normalized_cache=retained,
+        )
+
+    def training_table(self, method: str) -> np.ndarray:
+        """Return normalized training features, recomputing only if not retained."""
+        if method in self.normalized_cache:
+            return self.normalized_cache[method]
+        return self.normalizers[method].transform(self.features)
+
+    def targets_for(self, member: EnsembleMember) -> np.ndarray:
+        """Map original encoded class IDs to this member's model output IDs."""
+        if member.class_order is None:
+            return self.targets
+        return np.asarray(member.class_order)[self.targets]
+
+    def members_by_normalization(self) -> tuple[EnsembleMember, ...]:
+        """Return the stable execution order, also used when storing K/V caches.
+
+        Grouping lets members reuse a normalized training table. For eight
+        members this returns original indices 0, 4, 1, 5, 2, 6, 3, 7.
+        """
+        return tuple(
+            member for method in self.normalizers for member in self.members if member.normalization == method
+        )

causilo/data/encoding.py

@@ -0,0 +1,106 @@
+"""Learn a feature schema once and reuse its categories and column selection."""
+
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+from sklearn.preprocessing import OrdinalEncoder
+
+
+@dataclass
+class FeatureEncoder:
+    """Persist the input schema and ordinal encoder.
+
+    ``categories`` and ``continuous`` are input-column indices, not category
+    values. Encoded columns are ordered categorical first, then continuous;
+    ``retained`` is a varying-feature mask in that encoded order. These field
+    names are retained for compatibility with saved fitted states.
+    """
+
+    names: tuple | None
+    width: int
+    categories: tuple[int, ...]
+    continuous: tuple[int, ...]
+    encoder: OrdinalEncoder | None
+    retained: np.ndarray
+
+    @classmethod
+    def fit(cls, table) -> tuple["FeatureEncoder", np.ndarray]:
+        """Infer column kinds and return the schema plus nonconstant encoded features."""
+        is_dataframe = isinstance(table, pd.DataFrame)
+        array = np.asarray(table)
+        if array.ndim != 2 or min(array.shape) == 0:
+            raise ValueError("Features must be a nonempty two-dimensional table")
+        categorical_indices = []
+        for index in range(array.shape[1]):
+            dtype = table.dtypes.iloc[index] if is_dataframe else array.dtype
+            # Object arrays share one dtype; infer each column without treating
+            # numeric-looking strings as continuous measurements.
+            if not is_dataframe and pd.api.types.is_object_dtype(dtype):
+                inferred = pd.api.types.infer_dtype(array[:, index], skipna=True)
+                if inferred in {"integer", "floating", "mixed-integer-float", "decimal"}:
+                    continue
+            if pd.api.types.is_bool_dtype(dtype) or not pd.api.types.is_numeric_dtype(dtype):
+                if pd.api.types.is_datetime64_any_dtype(dtype) or pd.api.types.is_timedelta64_dtype(dtype):
+                    raise TypeError("Date and duration features must be encoded before fitting")
+                categorical_indices.append(index)
+        continuous_indices = tuple(i for i in range(array.shape[1]) if i not in categorical_indices)
+        schema = cls(
+            names=tuple(table.columns) if is_dataframe else None,
+            width=array.shape[1],
+            categories=tuple(categorical_indices),
+            continuous=continuous_indices,
+            encoder=None,
+            retained=np.ones(array.shape[1], dtype=bool),
+        )
+        if categorical_indices:
+            schema.encoder = OrdinalEncoder(
+                handle_unknown="use_encoded_value",
+                unknown_value=np.nan,
+                encoded_missing_value=np.nan,
+                dtype=np.float64,
+            )
+            schema.encoder.fit(schema._categorical(table))
+        numeric = schema.encode(table)
+        if len(numeric) > 1:
+            schema.retained = np.array([np.unique(column).size > 1 for column in numeric.T])
+        if not schema.retained.any():
+            raise ValueError("At least one varying feature is required")
+        return schema, numeric[:, schema.retained]
+
+    def _categorical(self, table) -> np.ndarray:
+        """Unify pandas/NumPy missing markers before ordinal encoding."""
+        values = np.asarray(table, dtype=object)[:, self.categories].copy()
+        values[pd.isna(values)] = np.nan
+        return values
+
+    def encode(self, table) -> np.ndarray:
+        """Validate the input schema and return categorical-first numeric columns.
+
+        Unknown categories become NaN so the model's missing-value embedding
+        handles them. Constant-column removal is applied by ``transform``.
+        """
+        if isinstance(table, pd.DataFrame):
+            if self.names is not None and tuple(table.columns) != self.names:
+                raise ValueError("Prediction columns must match fit columns in name and order")
+            numerical = table.iloc[:, list(self.continuous)].to_numpy(dtype=np.float64, na_value=np.nan)
+        else:
+            values = np.asarray(table)
+            if values.ndim != 2 or values.shape[1] != self.width:
+                raise ValueError(f"Expected {self.width} feature columns")
+            numerical = values[:, self.continuous]
+            if pd.api.types.is_object_dtype(numerical.dtype):
+                numerical = np.where(pd.isna(numerical), np.nan, numerical)
+            numerical = np.asarray(numerical, dtype=np.float64)
+        if np.asarray(table).shape[1] != self.width:
+            raise ValueError(f"Expected {self.width} feature columns")
+        parts = [self.encoder.transform(self._categorical(table))] if self.encoder is not None else []
+        parts.append(numerical)
+        result = np.concatenate(parts, axis=1)
+        if np.isinf(result).any():
+            raise ValueError("Infinite feature values are not supported")
+        return result
+
+    def transform(self, table) -> np.ndarray:
+        """Apply the fitted mappings and varying-feature mask without refitting."""
+        return self.encode(table)[:, self.retained]

causilo/data/ensemble.py

@@ -0,0 +1,87 @@
+"""Deterministic input permutations; all ensemble members share model weights."""
+
+import math
+import random
+from dataclasses import dataclass
+
+METHOD_ORDER = ("none", "rank2gaussian", "robust", "power")
+
+
+def _circular_distances(step: int, width: int) -> tuple[int, ...]:
+    """Distances to the first two neighbors when walking a ring with this stride."""
+    return tuple(
+        min((multiple * step) % width, (-multiple * step) % width)
+        for multiple in (1, 2)
+    )
+
+
+def make_affine_orders(width: int, count: int, random_state: int) -> list[tuple[int, ...]]:
+    """Build feature permutations while discouraging repeated local neighborhoods.
+
+    Coprime strides visit every feature exactly once. Among at most 32 candidate
+    strides, prefer the least-used distances to the first two ring neighbors.
+    Keep the first member in identity order. Separate local RNG streams preserve
+    the seeded policy without changing Python's global random state.
+    """
+    if width < 2:
+        return [tuple(range(width)) for _ in range(count)]
+    generator = random.Random(random_state)
+    ring = random.Random(random_state).sample(range(width), width)
+    strides = [step for step in range(1, width) if math.gcd(step, width) == 1]
+    usage = {}
+    orders = []
+    for index in range(count):
+        candidates = generator.sample(strides, min(32, len(strides))) if index else [1]
+        step = min(
+            candidates,
+            key=lambda stride: sum(usage.get(d, 0) for d in _circular_distances(stride, width)),
+        )
+        offset = generator.randrange(width) if index else 0
+        if index == 0:
+            order = tuple(range(width))
+        else:
+            order = tuple(ring[(offset + step * column) % width] for column in range(width))
+        orders.append(order)
+        for distance in _circular_distances(step, width):
+            usage[distance] = usage.get(distance, 0) + 1
+    return orders
+
+
+@dataclass(frozen=True)
+class EnsembleMember:
+    """One normalized/permuted input view of the shared model.
+
+    ``feature_order`` indexes the encoded, nonconstant feature table.
+    ``class_order[original_id]`` is the permuted class ID used as a target and
+    the model output column to read when restoring original class order.
+    Regression members have no class permutation.
+    """
+
+    normalization: str
+    feature_order: tuple[int, ...]
+    class_order: tuple[int, ...] | None
+
+
+def make_ensemble_members(
+    width: int, classes: int, count: int, random_state: int
+) -> tuple[EnsembleMember, ...]:
+    """Cycle normalizers and permutations in construction order, fixed by the seed."""
+    if count == 1:
+        return (EnsembleMember("none", tuple(range(width)), tuple(range(classes)) if classes else None),)
+    label_generator = random.Random(random_state)
+    class_orders = []
+    for index in range(count):
+        if not classes:
+            class_orders.append(None)
+            continue
+        if index % classes == 0:
+            # Rotate each sampled class ordering so labels occupy every slot
+            # before another base ordering is drawn.
+            base = label_generator.sample(range(classes), classes)
+        offset = index % classes
+        class_orders.append(tuple(base[(label + offset) % classes] for label in range(classes)))
+    feature_orders = make_affine_orders(width, count, random_state)
+    return tuple(
+        EnsembleMember(METHOD_ORDER[i % 4], feature_orders[i], class_orders[i])
+        for i in range(count)
+    )

causilo/data/normalization.py

@@ -0,0 +1,123 @@
+"""Training-fitted transforms with explicit missing masks and compressed outliers."""
+
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.special import ndtri
+from sklearn.preprocessing import PowerTransformer, RobustScaler, StandardScaler
+
+
+@dataclass
+class Rank2Gaussian:
+    """Interpolate training mid-ranks and map them through the Gaussian quantile function."""
+
+    knots: tuple[tuple[np.ndarray, np.ndarray], ...]
+
+    @classmethod
+    def fit(cls, table):
+        """Store unique values and tie-aware probabilities per feature column."""
+        knots = []
+        for column in table.T:
+            values, counts = np.unique(column, return_counts=True)
+            probabilities = (np.cumsum(counts) - counts / 2) / len(column)
+            knots.append((values, probabilities))
+        return cls(tuple(knots))
+
+    def transform(self, table):
+        """Reuse fitted ranks; interpolation clamps unseen extremes to endpoint ranks."""
+        return np.column_stack([
+            ndtri(np.interp(column, values, probabilities))
+            for column, (values, probabilities) in zip(table.T, self.knots)
+        ])
+
+
+@dataclass
+class Normalizer:
+    """Normalize encoded features using statistics fitted only on training rows.
+
+    Missing entries are temporarily mean-filled to fit/apply transforms, then
+    restored to NaN for the model's missing-value embedding. Method ``none``
+    skips the optional distribution transform, but still standardizes columns
+    and compresses outlier tails. Rank-to-Gaussian uses its own final scaling.
+    """
+
+    fill: np.ndarray
+    center: np.ndarray
+    spread: np.ndarray
+    normalizer: object | None
+    quantile_scale: StandardScaler | None
+    lower: np.ndarray
+    upper: np.ndarray
+
+    @classmethod
+    def fit(cls, table: np.ndarray, method: str) -> "Normalizer":
+        """Fit filling, scaling, optional distribution transform, and tail bounds."""
+        missing = np.isnan(table)
+        counts = (~missing).sum(axis=0)
+        fill = np.divide(
+            np.where(missing, 0, table).sum(axis=0), counts, out=np.zeros(table.shape[1]), where=counts > 0
+        )
+        filled = np.where(missing, fill, table)
+        if method == "rank2gaussian":
+            center = np.zeros(filled.shape[1], dtype=filled.dtype)
+            spread = np.ones_like(center)
+        else:
+            center, spread = filled.mean(axis=0), filled.std(axis=0)
+            tolerance = np.finfo(filled.dtype).eps * np.maximum(np.abs(center), 1)
+            spread = np.where(spread <= tolerance, 1.0, spread)
+        scaled = (filled - center) / spread
+        normalizer, quantile_scale = None, None
+        if method == "power":
+            normalizer = PowerTransformer().fit(scaled)
+        elif method == "robust":
+            normalizer = RobustScaler(unit_variance=True).fit(scaled)
+        elif method == "rank2gaussian":
+            normalizer = Rank2Gaussian.fit(scaled)
+        elif method != "none":
+            raise ValueError(f"Unknown internal normalization: {method}")
+        normalized = scaled if normalizer is None else normalizer.transform(scaled)
+        if method == "rank2gaussian":
+            quantile_scale = StandardScaler().fit(normalized)
+            normalized = quantile_scale.transform(normalized)
+        lower, upper = outlier_bounds(normalized)
+        return cls(
+            fill=fill,
+            center=center,
+            spread=spread,
+            normalizer=normalizer,
+            quantile_scale=quantile_scale,
+            lower=lower,
+            upper=upper,
+        )
+
+    def transform(self, table: np.ndarray) -> np.ndarray:
+        """Return float32 features with the original NaN mask and unchanged fitted state."""
+        missing = np.isnan(table)
+        values = (np.where(missing, self.fill, table) - self.center) / self.spread
+        if self.normalizer is not None:
+            values = self.normalizer.transform(values)
+        if self.quantile_scale is not None:
+            values = self.quantile_scale.transform(values)
+        values = compress_tails(values, self.lower, self.upper)
+        return np.where(missing, np.nan, values).astype(np.float32)
+
+
+def outlier_bounds(values):
+    """Estimate four-sigma bounds after excluding preliminary extreme values.
+
+    These thresholds are part of the fixed preprocessing policy. A second pass
+    limits how much extreme training values can inflate the final bounds.
+    """
+    correction = int(len(values) > 1)
+    preliminary = np.maximum(values.std(axis=0, ddof=correction), 1e-6)
+    mean = values.mean(axis=0)
+    central = np.where(np.abs(values - mean) > 4 * preliminary, np.nan, values)
+    center = np.nan_to_num(np.nanmean(central, axis=0), nan=0)
+    spread = np.maximum(np.nan_to_num(np.nanstd(central, axis=0, ddof=correction), nan=1), 1e-6)
+    return center - 4 * spread, center + 4 * spread
+
+
+def compress_tails(values, lower, upper):
+    """Leave central values unchanged and smoothly compress both tails with arcsinh."""
+    values = np.where(values < lower, lower - np.arcsinh(lower - values), values)
+    return np.where(values > upper, upper + np.arcsinh(values - upper), values)

causilo/engine.py

@@ -0,0 +1,156 @@
+"""Connect fixed pretrained weights to fitted data, caches, and ensemble outputs."""
+
+from dataclasses import dataclass
+from numbers import Integral
+
+import numpy as np
+import torch
+from sklearn.exceptions import NotFittedError
+
+from . import checkpoints
+from .data.dataset import PreparedDataset
+from .data.ensemble import EnsembleMember
+from .execution.cached import cached_predictions
+from .execution.direct import direct_predictions
+from .execution.runner import ModelCache, ModelRunner
+
+
+def resolve_device(request: str) -> torch.device:
+    """Resolve auto placement once; explicit unavailable devices fail instead of falling back."""
+    if request == "auto":
+        request = "cuda:0" if torch.cuda.is_available() else "cpu"
+    device = torch.device(request)
+    if device.type not in {"cpu", "cuda"}:
+        raise ValueError("Causilo supports CPU or a single CUDA device")
+    if device.type == "cuda":
+        index = device.index if device.index is not None else 0
+        if not torch.cuda.is_available() or index >= torch.cuda.device_count():
+            raise ValueError(f"Requested CUDA device {index} is unavailable")
+        device = torch.device("cuda", index)
+    return device
+
+
+@dataclass
+class FitState:
+    """Successful fit only; caches follow ``dataset.members_by_normalization()``.
+
+    Keep the stored tuple layout stable for fitted-state restoration. Consumers
+    obtain member/cache pairs through ``cached_members`` rather than rebuilding
+    the positional association themselves.
+    """
+
+    dataset: PreparedDataset
+    caches: tuple[ModelCache, ...] | None
+
+    def cached_members(self) -> tuple[tuple[EnsembleMember, ModelCache], ...]:
+        """Pair every cache with its input permutation, rejecting incomplete state."""
+        if self.caches is None:
+            raise ValueError("This fit does not contain K/V caches")
+        members = self.dataset.members_by_normalization()
+        if len(members) != len(self.caches):
+            raise ValueError("Cache count does not match the fitted ensemble")
+        return tuple(zip(members, self.caches))
+
+
+class Engine:
+    """Own one task/device model and replace its training context on each fit."""
+
+    def __init__(self, task: str, device: str) -> None:
+        self.task = task
+        self.device_request = device
+        self.device = resolve_device(device)
+        self.model = None
+        self.state: FitState | None = None
+
+    def fit(
+        self,
+        table,
+        targets,
+        *,
+        n_estimators: int,
+        use_kv_cache: bool,
+        retain_preprocessing: bool,
+        random_state: int,
+    ) -> None:
+        """Prepare context without optimizing weights; a failed fit stays unfitted."""
+        self.state = None
+        if isinstance(n_estimators, bool) or not isinstance(n_estimators, Integral) or n_estimators < 1:
+            raise ValueError("n_estimators must be a positive integer")
+        if isinstance(random_state, bool) or not isinstance(random_state, Integral) or random_state < 0:
+            raise ValueError("random_state must be a nonnegative integer")
+        for name, value in (("use_kv_cache", use_kv_cache), ("retain_preprocessing", retain_preprocessing)):
+            if not isinstance(value, (bool, np.bool_)):
+                raise ValueError(f"{name} must be a Boolean")
+        if self.model is None:
+            self.model = checkpoints.load_pretrained_model(self.task).to(self.device)
+        prepared = PreparedDataset.prepare(
+            table,
+            targets,
+            task=self.task,
+            n_estimators=int(n_estimators),
+            retain_preprocessing=retain_preprocessing,
+            max_classes=self.model.config.outputs,
+            random_state=int(random_state),
+        )
+        caches = None
+        if use_kv_cache:
+            runner = ModelRunner(self.model)
+            collected = []
+            with torch.inference_mode():
+                for member in prepared.members_by_normalization():
+                    features = prepared.training_table(member.normalization)[:, member.feature_order]
+                    y = prepared.targets_for(member)
+                    collected.append(runner.build_cache(self._tensor(features), self._tensor(y)))
+            caches = tuple(collected)
+        # Publish only a complete context; partially built caches stay local.
+        self.state = FitState(dataset=prepared, caches=caches)
+
+    def _tensor(self, value: np.ndarray) -> torch.Tensor:
+        """Place one member on the device and prepend its singleton batch axis."""
+        return torch.as_tensor(
+            np.ascontiguousarray(value), device=self.device, dtype=torch.float32
+        ).unsqueeze(0)
+
+    def require_state(self) -> FitState:
+        """Expose fitted state only after all preparation has succeeded."""
+        if self.state is None:
+            raise NotFittedError("Call fit successfully before prediction")
+        return self.state
+
+    def predict(self, table) -> np.ndarray:
+        """Return probabilities (queries, classes) or regression points (queries,)."""
+        state = self.require_state()
+        fitted = state.dataset
+        numeric = fitted.encoder.transform(table)
+        outputs = []
+
+        def reduce_output(result):
+            # Preserve the established sorted reduction order for regression.
+            # Sorting is algebraically unnecessary, but changing summation order
+            # can change floating-point results and saved-state reproducibility.
+            return result if self.task == "classification" else result.sort(dim=-1).values.mean(dim=-1)
+
+        with torch.inference_mode():
+            transformed = {
+                name: transform.transform(numeric) for name, transform in fitted.normalizers.items()
+            }
+            if state.caches is None:
+                predictions = direct_predictions(self.model, fitted, transformed, self.device, reduce_output)
+            else:
+                predictions = cached_predictions(
+                    self.model, fitted, transformed, state.cached_members(), self.device, reduce_output
+                )
+            for member, result in predictions:
+                if self.task == "classification":
+                    # class_order[original_id] gives the member's output column.
+                    outputs.append(result[:, member.class_order].numpy())
+                else:
+                    point = result.numpy().astype(np.float64)
+                    outputs.append(fitted.target_encoder.inverse_transform(point[:, None])[:, 0])
+        combined = np.mean(outputs, axis=0)
+        if self.task == "classification":
+            # Ensemble logits before softmax; averaging probabilities would
+            # implement a different prediction policy.
+            probabilities = np.exp(combined - combined.max(axis=-1, keepdims=True))
+            return probabilities / probabilities.sum(axis=-1, keepdims=True)
+        return combined