codadapt 1.0.0__py3-none-any.whl

codadapt/__init__.py

@@ -0,0 +1,8 @@
+"""CodAdapt: experimental adaptive integer addressing for tabular data."""
+
+from ._version import __version__
+from .classifier import CodAdaptClassifier
+from .regressor import CodAdaptRegressor
+
+CodAdapt = CodAdaptClassifier
+__all__ = ["CodAdapt", "CodAdaptClassifier", "CodAdaptRegressor", "__version__"]

codadapt/_base.py

@@ -0,0 +1,380 @@
+"""Shared estimator API and train-only preparation."""
+
+import copy
+import numbers
+import warnings
+from time import perf_counter
+
+import numpy as np
+import pandas as pd
+from scipy import sparse
+from sklearn.base import BaseEstimator
+from sklearn.exceptions import DataConversionWarning
+from sklearn.model_selection import train_test_split
+from sklearn.utils.multiclass import type_of_target
+from sklearn.utils.validation import check_is_fitted
+
+from ._core import addresses
+from ._training import (
+    predict_coded_memory_buckets,
+    train_coded_memory,
+)
+from .preprocessing import BucketEncoder
+
+
+def _rows(X, indices):
+    return X.iloc[indices] if isinstance(X, pd.DataFrame) else X[indices]
+
+
+def _target(y):
+    if y is None:
+        raise ValueError("This estimator requires y to be passed, but the target y is None.")
+    result = np.asarray(y)
+    if result.ndim == 2 and result.shape[1] == 1:
+        warnings.warn(
+            "A column-vector y was passed when a 1d array was expected.",
+            DataConversionWarning,
+            stacklevel=3,
+        )
+        result = result.ravel()
+    if result.ndim != 1:
+        raise ValueError("Only single-target, one-dimensional y is supported; no multioutput.")
+    if len(result) == 0 or pd.isna(result).any():
+        raise ValueError("Target y must be nonempty and contain no missing values (NaN or NA).")
+    if result.dtype.kind == "c":
+        raise ValueError("Complex data not supported in target y.")
+    if result.dtype.kind in "fiu" and not np.isfinite(result).all():
+        raise ValueError("Target y contains infinity or a value too large.")
+    return result
+
+
+class _BaseCodAdapt(BaseEstimator):
+    """Common parameters; construction only stores the supplied values."""
+
+    def __init__(
+        self,
+        *,
+        n_bins=32,
+        max_categories=128,
+        quantile_sample_size=20000,
+        n_tables=8,
+        features_per_table=3,
+        table_size=256,
+        main_effects=True,
+        l2=5.0,
+        categorical_features=None,
+        early_stopping=True,
+        validation_fraction=0.15,
+        tol=1e-5,
+        random_state=None,
+        verbosity=0,
+        lookup_budget=24,
+        memory_budget_mb=512,
+    ):
+        self.n_bins = n_bins
+        self.max_categories = max_categories
+        self.quantile_sample_size = quantile_sample_size
+        self.n_tables = n_tables
+        self.features_per_table = features_per_table
+        self.table_size = table_size
+        self.main_effects = main_effects
+        self.l2 = l2
+        self.categorical_features = categorical_features
+        self.early_stopping = early_stopping
+        self.validation_fraction = validation_fraction
+        self.tol = tol
+        self.random_state = random_state
+        self.verbosity = verbosity
+        self.lookup_budget = lookup_budget
+        self.memory_budget_mb = memory_budget_mb
+
+    def _validate_parameters(self):
+        if (
+            isinstance(self.memory_budget_mb, (bool, np.bool_))
+            or not isinstance(self.memory_budget_mb, numbers.Integral)
+            or self.memory_budget_mb < 1
+        ):
+            raise ValueError("memory_budget_mb must be a positive integer in MiB.")
+        bounds = {
+            "n_bins": (2, 4096),
+            "max_categories": (1, 4096),
+            "quantile_sample_size": (2, 1000000),
+            "n_tables": (0, 128),
+            "features_per_table": (1, 32),
+            "table_size": (16, 4096),
+            "verbosity": (0, 2),
+            "lookup_budget": (1, 4096),
+        }
+        for name, (low, high) in bounds.items():
+            value = getattr(self, name)
+            if isinstance(value, (bool, np.bool_)) or not isinstance(value, numbers.Integral):
+                raise ValueError(f"{name} must be an integer in [{low}, {high}].")
+            if not low <= value <= high:
+                raise ValueError(f"{name} must be in [{low}, {high}].")
+        if self.table_size & (self.table_size - 1):
+            raise ValueError("table_size must be a power of two from 16 through 4096.")
+        for name in ("main_effects", "early_stopping"):
+            if not isinstance(getattr(self, name), (bool, np.bool_)):
+                raise ValueError(f"{name} must be boolean.")
+        for name in ("l2", "tol", "validation_fraction"):
+            value = getattr(self, name)
+            if isinstance(value, (bool, np.bool_)) or not isinstance(value, numbers.Real):
+                raise ValueError(f"{name} must be a finite real number.")
+            if not np.isfinite(value):
+                raise ValueError(f"{name} must be finite.")
+        if self.l2 <= 0 or self.tol < 0:
+            raise ValueError("l2 must be positive and tol nonnegative.")
+        if not 0 < self.validation_fraction < 1:
+            raise ValueError("validation_fraction must lie strictly between zero and one.")
+        state = self.random_state
+        if state is not None and not isinstance(state, (numbers.Integral, np.random.RandomState)):
+            raise ValueError("random_state must be None, an integer, or a NumPy RandomState.")
+        if isinstance(state, numbers.Integral) and not 0 <= state <= 2**32 - 1:
+            raise ValueError("random_state integer must lie in [0, 2**32 - 1].")
+
+    def _check_memory_budget(self, n, p, eval_set=None):
+        """Apply the deterministic working-allocation heuristic, not an RSS cap."""
+        validation_rows = 0
+        if eval_set is not None:
+            if not isinstance(eval_set, (tuple, list)) or len(eval_set) != 2:
+                raise ValueError("eval_set must be the pair (X_valid, y_valid).")
+            if sparse.issparse(eval_set[0]):
+                raise TypeError(
+                    "Sparse validation input is not supported; provide a dense numeric array."
+                )
+            # Read only the row count before copies/encoding; detailed validation
+            # remains in preparation. External validation is live during fit too.
+            try:
+                validation_rows = len(eval_set[0])
+            except TypeError as error:
+                raise ValueError("Validation X must be a 2D array or DataFrame.") from error
+        estimate = (int(n) + int(validation_rows)) * (40 * int(p) + 2 * self.n_tables + 80)
+        estimate += (
+            24
+            * (p + self.n_tables * min(p, self.features_per_table))
+            * max(self.n_bins + 1, self.max_categories + 3)
+        )
+        estimate += 32 * self.n_tables * self.table_size
+        budget = int(self.memory_budget_mb) * 1024**2
+        if estimate > budget:
+            raise ValueError(
+                f"Estimated working allocation {estimate / 1024**2:.2f} MiB exceeds "
+                f"the {self.memory_budget_mb} MiB memory_budget_mb budget "
+                f"({n} training rows, {validation_rows} external validation rows, "
+                f"{p} features). Use fewer rows/features or explicitly raise "
+                "memory_budget_mb only with sufficient RAM headroom. This heuristic "
+                "is not a total process RSS limit or an out-of-memory guarantee."
+            )
+        return int(estimate)
+
+    def __setstate__(self, state):
+        # Trusted RC1 artifacts predate the explicit resource-only parameter.
+        state = state.copy()
+        state.setdefault("memory_budget_mb", 512)
+        super().__setstate__(state)
+
+    def _encode_target(self, y, fitting=False):
+        if self._task == "regression":
+            try:
+                values = np.asarray(y, dtype=np.float64)
+            except (TypeError, ValueError) as error:
+                raise ValueError("Regression requires a finite numeric target.") from error
+            if not np.isfinite(values).all():
+                raise ValueError("Regression target must be finite.")
+            return values
+        if fitting:
+            try:
+                target_type = type_of_target(y)
+                classes = np.unique(y)
+            except (TypeError, ValueError) as error:
+                raise ValueError(
+                    "Binary targets require comparable scalar class labels."
+                ) from error
+            if target_type == "continuous":
+                raise ValueError(
+                    "Unknown label type: continuous. Only binary classification is supported."
+                )
+            if target_type != "binary" or len(classes) != 2:
+                raise ValueError(
+                    "Only binary classification is supported: one class was found; "
+                    "need two classes with positive weight."
+                )
+            self.classes_ = classes
+        if not np.isin(y, self.classes_).all():
+            raise ValueError(
+                "Validation target contains classes absent from positive-weight training."
+            )
+        return (y == self.classes_[1]).astype(np.float64)
+
+    def fit(self, X, y, sample_weight=None, eval_set=None):
+        """Fit only on positive-weight training rows, with optional external validation."""
+        started = perf_counter()
+        # A failed refit must not leave the former fitted model available for prediction.
+        for name in list(vars(self)):
+            if name.endswith("_"):
+                delattr(self, name)
+        self._validate_parameters()
+        if sparse.issparse(X):
+            raise TypeError("Sparse input is not supported; provide a dense numeric array.")
+        if not isinstance(X, pd.DataFrame):
+            try:
+                X = np.asarray(X)
+            except (TypeError, ValueError) as error:
+                raise ValueError(
+                    "X must be a DataFrame or dense numeric array-like input."
+                ) from error
+        if X.ndim != 2:
+            raise ValueError("Expected 2D array. Reshape your data to (n_samples, n_features).")
+        n, p = X.shape
+        if not p:
+            raise ValueError(
+                f"Found array with 0 feature(s) (shape={X.shape}) while a minimum of 1 is required."
+            )
+        if not n:
+            raise ValueError(
+                f"Found array with 0 sample(s) (shape={X.shape}); a minimum of 1 is required."
+            )
+        y = _target(y)
+        if len(y) != n:
+            raise ValueError("X and y have inconsistent numbers of samples.")
+        if sample_weight is None:
+            weight = np.ones(n, dtype=np.float64)
+        else:
+            try:
+                weight = np.asarray(sample_weight, dtype=np.float64)
+            except (TypeError, ValueError) as error:
+                raise ValueError("sample_weight must be a finite nonnegative 1D array.") from error
+            if weight.ndim != 1 or len(weight) != n:
+                raise ValueError("sample_weight must be a 1D array with one value per sample.")
+            if not np.isfinite(weight).all() or (weight < 0).any():
+                raise ValueError("sample_weight must be finite and nonnegative.")
+        if not np.isfinite(weight.sum()) or weight.sum() <= 0:
+            raise ValueError(
+                "sample_weight must contain at least one non-zero value and have a finite positive sum."
+            )
+        self.estimated_memory_bytes_ = self._check_memory_budget(n, p, eval_set)
+        keep = np.flatnonzero(weight > 0)
+        X_train, y_train, weight = _rows(X, keep), y[keep], weight[keep].copy()
+        y_train = self._encode_target(y_train, fitting=True)
+        rng = (
+            copy.deepcopy(self.random_state)
+            if isinstance(self.random_state, np.random.RandomState)
+            else np.random.RandomState(self.random_state)
+        )
+        X_valid = y_valid = None
+        if eval_set is not None:
+            if not isinstance(eval_set, (tuple, list)) or len(eval_set) != 2:
+                raise ValueError("eval_set must be the pair (X_valid, y_valid).")
+            X_valid, y_valid = eval_set
+            if sparse.issparse(X_valid):
+                raise TypeError(
+                    "Sparse validation input is not supported; provide a dense numeric array."
+                )
+            if not isinstance(X_valid, pd.DataFrame):
+                try:
+                    X_valid = np.asarray(X_valid)
+                except (TypeError, ValueError) as error:
+                    raise ValueError(
+                        "Validation X must be a 2D dense array or DataFrame."
+                    ) from error
+            if X_valid.ndim != 2:
+                raise ValueError(
+                    "Validation X must be a 2D array with shape (n_samples, n_features)."
+                )
+            y_valid = self._encode_target(_target(y_valid))
+            if len(y_valid) != X_valid.shape[0]:
+                raise ValueError("Validation X and y have inconsistent numbers of samples.")
+        elif self.early_stopping:
+            indices = np.arange(len(y_train))
+            try:
+                train_indices, valid_indices = train_test_split(
+                    indices,
+                    test_size=self.validation_fraction,
+                    random_state=copy.deepcopy(rng),
+                    stratify=y_train if self._task == "classification" else None,
+                )
+                if self._task == "classification" and len(np.unique(y_train[train_indices])) != 2:
+                    raise ValueError("Training split does not retain both classes.")
+            except ValueError as error:
+                raise ValueError(
+                    f"Insufficient data ({len(y_train)} sample(s)) for internal validation; "
+                    "use early_stopping=False or a valid eval_set."
+                ) from error
+            X_valid, y_valid = _rows(X_train, valid_indices), y_train[valid_indices]
+            X_train, y_train, weight = (
+                _rows(X_train, train_indices),
+                y_train[train_indices],
+                weight[train_indices],
+            )
+        self.n_features_in_ = p
+        self.timings_ = {}
+        preprocessing_start = perf_counter()
+        encoder_seed = int(rng.randint(2**31 - 1))
+        finest = max(2, min(16, self.n_bins))
+        self.resolution_bins_ = [max(2, finest // 4), max(2, finest // 2), finest]
+        self.shared_encoder_ = BucketEncoder(
+            n_bins=finest,
+            max_categories=self.max_categories,
+            quantile_sample_size=self.quantile_sample_size,
+            categorical_features=self.categorical_features,
+            random_state=encoder_seed,
+        )
+        Z_levels = self.shared_encoder_.fit_transform_levels(X_train, self.resolution_bins_)
+        Z_valid_levels = (
+            self.shared_encoder_.transform_levels(X_valid) if X_valid is not None else None
+        )
+        self.encoders_ = self.shared_encoder_.level_views_
+        self.encoder_ = self.shared_encoder_
+        if hasattr(self.encoder_, "feature_names_in_"):
+            self.feature_names_in_ = self.encoder_.feature_names_in_.copy()
+        self.timings_["preprocessing"] = perf_counter() - preprocessing_start
+        train_coded_memory(
+            self,
+            Z_levels,
+            y_train,
+            weight,
+            Z_valid_levels,
+            y_valid,
+            rng,
+        )
+        self.timings_["fit_total"] = perf_counter() - started
+        if self.verbosity:
+            print(
+                f"CodAdapt fitted {self.n_iter_} iterations; restored={self.best_iteration_}; moves={self.accepted_moves_}; fit={self.timings_['fit_total']:.3f}s"
+            )
+        self._is_fitted_ = True
+        return self
+
+    def _raw_predict_buckets(self, Z):
+        result = np.full(len(Z), self.intercept_, dtype=np.float64)
+        for j, values in enumerate(self.main_tables_):
+            result += values[Z[:, j]]
+        for features, codes, values in zip(self.groups_, self.codes_, self.tables_):
+            index = addresses(Z, features, codes, self.table_size)
+            supported = np.ones(len(Z), dtype=bool)
+            for j in features:
+                supported &= self.encoder_.support_[j][Z[:, j]]
+            result += np.where(supported, values[index], 0.0)
+        return result
+
+    def _decision(self, X):
+        check_is_fitted(self)
+        Z_levels = self.shared_encoder_.transform_levels(X)
+        return predict_coded_memory_buckets(
+            self.levels_, self.encoders_, Z_levels, self.intercept_, self.table_size
+        )
+
+    def __sklearn_is_fitted__(self):
+        return getattr(self, "_is_fitted_", False)
+
+    def _more_tags(self):
+        return {"allow_nan": True, "requires_y": True, "X_types": ["2darray", "categorical"]}
+
+    def __sklearn_tags__(self):
+        tags = super().__sklearn_tags__()
+        tags.input_tags.allow_nan = True
+        tags.input_tags.categorical = True
+        tags.input_tags.string = True
+        tags.input_tags.sparse = False
+        tags.target_tags.required = True
+        return tags

codadapt/_core.py

@@ -0,0 +1,70 @@
+"""Vectorized addressing and exact incremental sufficient statistics.
+
+The candidate loop only touches arrays of length ``table_size``. Row indices
+are grouped once by training; each visited coordinate aggregates its rows once.
+"""
+
+import numpy as np
+
+GAIN_RTOL = 1e-12
+_CANCELLATION_RTOL = 64 * np.finfo(np.float64).eps
+
+
+def sigmoid(scores):
+    """Stable logistic function, including scores whose exponent would overflow."""
+    scores = np.asarray(scores, dtype=np.float64)
+    result = np.empty_like(scores)
+    positive = scores >= 0
+    result[positive] = 1.0 / (1.0 + np.exp(-scores[positive]))
+    exponent = np.exp(scores[~positive])
+    result[~positive] = exponent / (1.0 + exponent)
+    return result
+
+
+def addresses(z, features, codes, table_size):
+    """Return compact addresses; ``codes[k]`` belongs to global ``features[k]``.
+
+    Add in int64 before the modulo: compact input code dtypes must not determine
+    arithmetic precision. Parameter and feature validation belongs to the API.
+    """
+    total = np.zeros(z.shape[0], dtype=np.int64)
+    for feature, code in zip(features, codes):
+        total += code[z[:, feature]].astype(np.int64, copy=False)
+    return (total % table_size).astype(np.uint16)
+
+
+def aggregate(address, h, r, table_size):
+    """Aggregate mass A and weighted response B in float64 for one addressing."""
+    h, r = np.asarray(h, dtype=np.float64), np.asarray(r, dtype=np.float64)
+    a = np.bincount(address, weights=h, minlength=table_size)
+    b = np.bincount(address, weights=h * r, minlength=table_size)
+    return a, b
+
+
+def gain(a, b, l2):
+    """The maximized regularized quadratic improvement, not exact logistic gain."""
+    result = 0.5 * np.dot(b, b / (a + l2))
+    if not np.isfinite(result):
+        raise FloatingPointError("Non-finite quadratic gain; rescale target or sample_weight.")
+    return float(result)
+
+
+def incremental_stats(a, b, d, e, delta):
+    """Move a group's statistics forward by delta cells, without mutating input.
+
+    D/E refer to the current addresses of the group. Negative residual mass is
+    clipped only within 64 machine epsilons of the per-cell arithmetic scale;
+    larger negative mass signals a broken invariant and raises an exception.
+    Zero-mass cancellation in B is corrected only at the same relative scale.
+    """
+    a, b, d, e = (np.asarray(value, dtype=np.float64) for value in (a, b, d, e))
+    remaining_a, remaining_b = a - d, b - e
+    mass_tolerance = _CANCELLATION_RTOL * np.maximum(1.0, np.maximum(np.abs(a), np.abs(d)))
+    if np.any(remaining_a < -mass_tolerance):
+        raise FloatingPointError("Significant negative mass in incremental statistics.")
+    remaining_a[remaining_a < 0] = 0.0
+    response_tolerance = _CANCELLATION_RTOL * np.maximum(1.0, np.maximum(np.abs(b), np.abs(e)))
+    cancelled = (remaining_a == 0) & (np.abs(remaining_b) <= response_tolerance)
+    remaining_b[cancelled] = 0.0
+    shift = int(delta) % len(a)
+    return remaining_a + np.roll(d, shift), remaining_b + np.roll(e, shift)

codadapt/_multiclass.py

@@ -0,0 +1,229 @@
+"""Exact native multiclass preparation and fit-local X-only reuse."""
+
+import copy
+import warnings
+from dataclasses import dataclass
+from time import perf_counter
+
+import numpy as np
+import pandas as pd
+from scipy import sparse
+from sklearn.model_selection import train_test_split
+from sklearn.utils.multiclass import type_of_target
+
+from ._base import _rows, _target
+from ._core import addresses, aggregate
+from ._training import _surrogate, train_coded_memory
+from .preprocessing import BucketEncoder
+
+
+def validate_weights(weight, n):
+    if weight is None:
+        result = np.ones(n, dtype=np.float64)
+    else:
+        try:
+            result = np.asarray(weight, dtype=np.float64)
+        except (TypeError, ValueError) as error:
+            raise ValueError("sample_weight must be a finite nonnegative 1D array.") from error
+        if result.ndim != 1 or len(result) != n:
+            raise ValueError("sample_weight must be a 1D array with one value per sample.")
+        if not np.isfinite(result).all() or (result < 0).any():
+            raise ValueError("sample_weight must be finite and nonnegative.")
+    if not np.isfinite(result.sum()) or result.sum() <= 0:
+        raise ValueError(
+            "sample_weight must contain at least one non-zero value and have a finite positive sum."
+        )
+    return result
+
+
+def is_multiclass(y, weight):
+    """Binary/invalid cases still delegate to the original public fit verbatim."""
+    try:
+        a = np.asarray(y)
+        if a.ndim == 2 and a.shape[1] == 1:
+            a = a.ravel()
+        if a.ndim != 1 or not len(a) or pd.isna(a).any():
+            return False
+        if weight is not None:
+            w = np.asarray(weight, dtype=float)
+            if w.ndim == 1 and len(w) == len(a) and np.isfinite(w).all() and (w >= 0).all():
+                a = a[w > 0]
+        return type_of_target(a) == "multiclass" and len(np.unique(a)) > 2
+    except (TypeError, ValueError):
+        return False
+
+
+@dataclass
+class Prepared:
+    x: object
+    y: np.ndarray
+    weight: np.ndarray
+    xv: object
+    yv: object
+    classes: np.ndarray
+    rng: object
+    train_ids: np.ndarray
+    valid_ids: object
+    estimated_bytes: int
+
+
+def prepare(model, X, y, sample_weight, eval_set):
+    model._validate_parameters()
+    if sparse.issparse(X):
+        raise TypeError("Sparse input is not supported; provide a dense numeric array.")
+    if not isinstance(X, pd.DataFrame):
+        X = np.asarray(X)
+    if X.ndim != 2:
+        raise ValueError("Expected 2D array. Reshape your data to (n_samples, n_features).")
+    n, p = X.shape
+    if not n or not p:
+        raise ValueError(f"Found array with {n} sample(s) and {p} feature(s); minimum1 required.")
+    y = _target(y)
+    if len(y) != n:
+        raise ValueError("X and y have inconsistent numbers of samples.")
+    weight = validate_weights(sample_weight, n)
+    estimate = model._check_memory_budget(n, p, eval_set)
+    keep = np.flatnonzero(weight > 0)
+    x, y, w = _rows(X, keep), y[keep], weight[keep].copy()
+    try:
+        classes, counts = np.unique(y, return_counts=True)
+        kind = type_of_target(y)
+    except (TypeError, ValueError) as error:
+        raise ValueError("Classification requires comparable scalar class labels.") from error
+    if kind != "multiclass" or len(classes) < 3:
+        raise ValueError("At least three positive-weight classes required for shared OVR.")
+    rng = (
+        copy.deepcopy(model.random_state)
+        if isinstance(model.random_state, np.random.RandomState)
+        else np.random.RandomState(model.random_state)
+    )
+    xv = yv = valid_ids = None
+    train_ids = keep
+    if eval_set is not None:
+        if not isinstance(eval_set, (tuple, list)) or len(eval_set) != 2:
+            raise ValueError("eval_set must be the pair (X_valid, y_valid).")
+        xv, yv = eval_set
+        if sparse.issparse(xv):
+            raise TypeError(
+                "Sparse validation input is not supported; provide a dense numeric array."
+            )
+        if not isinstance(xv, pd.DataFrame):
+            xv = np.asarray(xv)
+        yv = _target(yv)
+        if xv.ndim != 2 or len(yv) != len(xv):
+            raise ValueError("Validation X and y have inconsistent numbers of samples.")
+        if not np.isin(yv, classes).all():
+            raise ValueError(
+                "Validation target contains classes absent from positive-weight training."
+            )
+    elif model.early_stopping:
+        try:
+            a, b = train_test_split(
+                np.arange(len(y)),
+                test_size=model.validation_fraction,
+                random_state=copy.deepcopy(rng),
+                stratify=y,
+            )
+            if len(np.unique(y[a])) != len(classes):
+                raise ValueError("Training split does not retain all classes.")
+        except ValueError as error:
+            raise ValueError(
+                f"Insufficient data ({len(y)} sample(s)) for common stratified internal validation; use early_stopping=False or a valid eval_set."
+            ) from error
+        xv, yv = _rows(x, b), y[b]
+        x, y, w = _rows(x, a), y[a], w[a]
+        train_ids, valid_ids = keep[a], keep[b]
+    _, counts = np.unique(y, return_counts=True)
+    if counts.min() < 5:
+        warnings.warn(
+            "Rare class (<5 positive-weight TRAIN examples): native OVR quality is not guaranteed.",
+            UserWarning,
+            stacklevel=2,
+        )
+    return Prepared(x, y, w, xv, yv, classes, rng, train_ids, valid_ids, int(estimate))
+
+
+@dataclass
+class NativeHead:
+    intercept: float
+    levels: list
+    table_size: int
+
+
+class XCache:
+    def __init__(self, weight, budget=32 * 2**20):
+        self.weight, self.budget, self.bytes = weight, budget, 0
+        self.addresses, self.masses, self.combinations = {}, {}, {}
+
+    def address(self, z, features, codes, size):
+        key = (id(z), tuple(features), size, tuple(c.tobytes() for c in codes))
+        if key not in self.addresses:
+            value = addresses(z, features, codes, size)
+            if self.bytes + value.nbytes <= self.budget:
+                self.addresses[key] = value
+                self.bytes += value.nbytes
+            return value
+        return self.addresses[key]
+
+    def surrogate(self, y, score, old, weight, binary):
+        assert binary and weight is self.weight
+        return _surrogate(y, score, old, weight, binary)
+
+    def aggregate(self, index, h, r, size):
+        # native h=weight/4 for EVERY class in this fit; weight is identity bound.
+        key = (index.__array_interface__["data"][0], index.shape, index.strides, size)
+        if key in self.masses:
+            return self.masses[key][1], np.bincount(index, weights=h * r, minlength=size)
+        a, b = aggregate(index, h, r, size)
+        required = a.nbytes + (0 if index.base is not None else index.nbytes)
+        if self.bytes + required <= self.budget:
+            self.masses[key] = (index, a)
+            self.bytes += required
+        return a, b
+
+    def combination_count(self, z, features):
+        key = (id(z), tuple(features))
+        if key not in self.combinations:
+            self.combinations[key] = int(np.unique(z[:, features], axis=0).shape[0])
+        return self.combinations[key]
+
+
+def fit_heads(model, d):
+    tick = perf_counter()
+    rng = copy.deepcopy(d.rng)
+    encoder_seed = int(rng.randint(2**31 - 1))
+    finest = max(2, min(16, model.n_bins))
+    resolutions = [max(2, finest // 4), max(2, finest // 2), finest]
+    encoder = BucketEncoder(
+        n_bins=finest,
+        max_categories=model.max_categories,
+        quantile_sample_size=model.quantile_sample_size,
+        categorical_features=model.categorical_features,
+        random_state=encoder_seed,
+    )
+    z = encoder.fit_transform_levels(d.x, resolutions)
+    zv = encoder.transform_levels(d.xv) if d.xv is not None else None
+    preprocessing_s = perf_counter() - tick
+    cache = XCache(d.weight)
+    heads = []
+    # Local native head models are discarded; only predictive state survives.
+    from codadapt import CodAdaptClassifier
+
+    for label in d.classes:
+        head = CodAdaptClassifier(**model.get_params(deep=False))
+        head.classes_ = np.array([0, 1])
+        head.shared_encoder_ = head.encoder_ = encoder
+        head.encoders_, head.resolution_bins_ = encoder.level_views_, resolutions
+        head.n_features_in_, head.timings_ = d.x.shape[1], {"preprocessing": 0.0}
+        train_coded_memory(
+            head,
+            z,
+            (d.y == label).astype(float),
+            d.weight,
+            zv,
+            (d.yv == label).astype(float) if d.yv is not None else None,
+            copy.deepcopy(rng),
+            cache=cache,
+        )
+        heads.append(NativeHead(head.intercept_, head.levels_, head.table_size))
+    return encoder, heads, preprocessing_s