fix(dml): fail closed when cross-fit learners cannot be isolated

changelog.d/20260903-dml-crossfit-learner-isolation.md

@@ -0,0 +1,7 @@
+### Behavioral Changes
+- **Cross-fit learner isolation now fails closed**: DMLDiD replaces the prior
+  warning-and-reuse fallback. It raises `TypeError` before fitting any
+  group-time cell when a custom learner template's `deepcopy` fails or returns
+  the original object. Implement `__deepcopy__` to return an independent
+  instance; custom implementations remain responsible for nested mutable state.
+  A failed re-fit now also clears any previous fitted result.

diff_diff/_crossfit.py

@@ -10,11 +10,12 @@
 
 ``cross_fit_predict`` produces out-of-fold nuisance predictions for EVERY
 unit: for each fold k the learner is fit on ``train_mask(k) & fit_mask`` and
-predicts all units in fold k. Each fold fits a DEEP COPY of the user's
-(never-fit) learner template, so no state — nested estimators and container
-parameters included — can carry across folds; an un-deep-copyable learner is
-reused with a loud warning under the fit-reset contract (see
-``diff_diff._learners``).
+predicts all units in fold k. Each fold uses a distinct top-level DEEP COPY of
+the user's (never-fit) learner template, so the template itself cannot carry
+fitted state across folds. Custom ``__deepcopy__`` implementations remain
+responsible for isolating nested mutable state. A template that cannot be
+deep-copied to a distinct top-level object fails closed with ``TypeError``
+(see ``diff_diff._learners``).
 
 Exception semantics (determinate):
 
@@ -32,8 +33,8 @@
 """
 
 import copy
+import inspect
 import pickle
-import warnings
 from dataclasses import dataclass, field
 from typing import Any, Dict, Iterator, Literal, Optional, Tuple, cast, overload
 
@@ -58,38 +59,68 @@
 _LOG_LOSS_CLIP = 1e-15
 
 
-def _fresh_learner(learner: Any) -> Any:
-    """Per-fold learner isolation: a deep copy of the (never-fit) template.
+def _clone_learner_template(learner: Any, *, label: str) -> Any:
+    """Return a distinct deep copy or raise a sanitised ``TypeError``.
 
-    ``copy.deepcopy`` of the user's template gives every fold a fully
-    independent learner — nested estimators, estimators inside lists/dicts,
-    accumulators, and warm-start state included — so no state (and therefore
-    no data from a previous complement, which includes the current evaluation
-    fold) can carry across folds. This is strictly stronger than
-    get_params-based reconstruction (which shares any estimator stored inside
-    a container parameter). The template itself is never fit. A copy FAILURE
-    is never silent: the instance is reused with a loud ``UserWarning`` naming
-    the learner and the fit-reset assumption now being relied on
-    (no-silent-failures rule).
+    The check deliberately proves only top-level identity. A custom
+    ``__deepcopy__`` remains responsible for isolating nested mutable state.
+    Error messages expose the learner class and copy exception class, never
+    foreign exception text that could carry credentials, paths, or data.
     """
     try:
-        return copy.deepcopy(learner)
-    except Exception as exc:  # noqa: BLE001 - loud fallback, never silent
-        # Exception CLASS only, never the message: a foreign learner's
-        # __deepcopy__ error text can embed credentials/paths/data excerpts,
-        # and this warning lands in notebook/CI logs (the same boundary as
-        # DMLDiD's persisted-diagnostics sanitization).
-        warnings.warn(
-            f"cross_fit_predict: could not deep-copy the "
-            f"{type(learner).__name__} template for this fold "
-            f"({type(exc).__name__}); "
-            "REUSING the same instance and relying on its fit-reset behavior. "
-            "A warm-start/stateful learner in this situation can leak data "
-            "across folds.",
-            UserWarning,
-            stacklevel=3,
+        clone = copy.deepcopy(learner)
+    except Exception as exc:  # noqa: BLE001 - sanitize a foreign exception boundary
+        copy_error_class = type(exc).__name__
+    else:
+        if clone is not learner:
+            return clone
+        raise TypeError(
+            f"{label}: {type(learner).__name__} learner template's __deepcopy__ "
+            "returned the original object; implement __deepcopy__ to return an "
+            "independent instance."
         )
-        return learner
+
+    raise TypeError(
+        f"{label}: could not deep-copy {type(learner).__name__} learner template "
+        f"({copy_error_class}); implement __deepcopy__ to return an independent instance."
+    )
+
+
+def _probe_learner_cloneability(learner: Any, *, kind: str, param_name: str) -> Any:
+    """Return an independent clone after validating its learner protocol."""
+    clone = _clone_learner_template(learner, label=param_name)
+    validate_learner(clone, kind=kind, param_name=param_name)
+    return clone
+
+
+def _validate_sample_weight_support(learner: Any, *, param_name: str) -> None:
+    """Require a learner ``fit`` method that accepts ``sample_weight`` by keyword."""
+    try:
+        sig = inspect.signature(learner.fit)
+    except (TypeError, ValueError):  # pragma: no cover - exotic callables
+        return  # Cannot introspect; let the learner surface any fit-time error.
+    for param in sig.parameters.values():
+        if param.kind is inspect.Parameter.VAR_KEYWORD:
+            return
+        if param.name == "sample_weight" and param.kind in (
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            inspect.Parameter.KEYWORD_ONLY,
+        ):
+            return
+    raise TypeError(
+        f"{param_name}: learner object {type(learner).__name__!r} must accept "
+        "sample_weight by keyword in fit(); add a sample_weight parameter (or **kwargs)."
+    )
+
+
+def _fresh_learner(learner: Any, *, context_label: str, fold: int) -> Any:
+    """Return an isolated learner for one fold; fail closed as a backstop.
+
+    DMLDiD probes user templates before fitting any cell. Direct callers of
+    ``cross_fit_predict`` still receive the same no-silent-failures contract.
+    """
+    label = f"{context_label}: fold {fold}" if context_label else f"cross_fit_predict: fold {fold}"
+    return _clone_learner_template(learner, label=label)
 
 
 def _unique_or_raise(arr: np.ndarray, name: str, **kwargs: Any) -> Any:
@@ -547,7 +578,10 @@ def cross_fit_predict(
 
         # (b) Learner errors during the fold -> DegenerateFoldError, chained.
         try:
-            fold_learner = _fresh_learner(learner)
+            fold_learner = _fresh_learner(learner, context_label=context_label, fold=k)
+            validate_learner(fold_learner, kind=kind, param_name=f"{label}fold {k} learner")
+            if w_fit is not None:
+                _validate_sample_weight_support(fold_learner, param_name=f"{label}fold {k} learner")
             # Unweighted path calls fit(X, y) WITHOUT the keyword: the
             # advertised duck-typed contract is fit/predict(_proba), so a
             # learner whose fit signature is only (X, y) must work when no

diff_diff/_learners.py

@@ -4,10 +4,10 @@
 ----------------
 Learners ALWAYS receive a raw covariate matrix ``X`` with NO intercept column;
 every learner manages the intercept internally (sklearn convention). Learners
-are INSTANCES with sklearn fit-reset semantics: ``fit`` fully re-initializes
-the fitted state and returns ``self``. A stateful/warm-start user learner that
-violates fit-reset cannot be detected without taking a clone dependency —
-documented accepted limitation.
+are INSTANCES whose ``fit`` returns ``self``. Cross-fitting requires each
+template to ``deepcopy`` to a distinct top-level object; a custom
+``__deepcopy__`` implementation is responsible for isolating nested mutable
+state.
 
 Native learners (``"linear"``, ``"ridge"``, ``"logit"``, ``"sieve"``) wrap
 the ``diff_diff.linalg`` solvers and expose sklearn-style fitted state

diff_diff/dml_did.py

@@ -32,7 +32,6 @@
 """
 
 import decimal
-import inspect
 import secrets
 import warnings
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union, cast
@@ -41,7 +40,13 @@
 import pandas as pd
 
 from diff_diff._base import BaseEstimator
-from diff_diff._crossfit import DegenerateFoldError, assign_folds, cross_fit_predict
+from diff_diff._crossfit import (
+    DegenerateFoldError,
+    _probe_learner_cloneability,
+    _validate_sample_weight_support,
+    assign_folds,
+    cross_fit_predict,
+)
 from diff_diff._dr_scores import (
     _chang_rcs_score_augmented_with_slope,
     chang_panel_score,
@@ -99,7 +104,9 @@ def _validate_n_folds(value: Any) -> int:
     return int(value)
 
 
-def _validate_learner_spec(spec: Any, *, kind: str, param_name: str) -> None:
+def _validate_learner_spec(
+    spec: Any, *, kind: str, param_name: str, require_sample_weight: bool = False
+) -> None:
     """Eager learner-spec validation naming the ACTUAL constructor param.
 
     ``make_learner`` hard-codes ``param_name="learner"`` for objects, which
@@ -114,38 +121,9 @@ def _validate_learner_spec(spec: Any, *, kind: str, param_name: str) -> None:
             )
         return
     validate_learner(spec, kind=kind, param_name=param_name)
-
-
-def _validate_learner_sample_weight_support(spec: Any, param_name: str) -> None:
-    """Reject a user learner whose ``fit`` cannot take ``sample_weight``.
-
-    Declared-survey fits pass ``sample_weight`` into ``cross_fit_predict``,
-    which forwards it BY KEYWORD (``fit_kwargs = {"sample_weight": w_fit}``)
-    and deliberately propagates the learner's ``TypeError`` — so a learner
-    whose ``fit`` has neither a keyword-addressable ``sample_weight``
-    parameter (POSITIONAL_OR_KEYWORD or KEYWORD_ONLY; POSITIONAL_ONLY does
-    not qualify) nor ``**kwargs`` would hard-crash mid-fit. Raises
-    ``TypeError`` up front instead (the ``validate_learner`` convention for
-    object-capability failures).
-    """
-    try:
-        sig = inspect.signature(spec.fit)
-    except (TypeError, ValueError):  # pragma: no cover - exotic callables
-        return  # cannot introspect; let cross_fit_predict surface any error
-    for param in sig.parameters.values():
-        if param.kind is inspect.Parameter.VAR_KEYWORD:
-            return
-        if param.name == "sample_weight" and param.kind in (
-            inspect.Parameter.POSITIONAL_OR_KEYWORD,
-            inspect.Parameter.KEYWORD_ONLY,
-        ):
-            return
-    raise TypeError(
-        f"survey_design= requires learners whose fit() accepts sample_weight "
-        f"by keyword; the {param_name} object {type(spec).__name__!r} does not. "
-        "Add a sample_weight parameter (or **kwargs) to its fit(), or use a "
-        "library-native learner name."
-    )
+    clone = _probe_learner_cloneability(spec, kind=kind, param_name=param_name)
+    if require_sample_weight:
+        _validate_sample_weight_support(clone, param_name=param_name)
 
 
 def _raw_label_is_infinite(value: Any) -> bool:
@@ -398,7 +376,7 @@ def __init__(
         self.results_: Optional[DMLDiDResults] = None
         self.is_fitted_ = False
 
-    def _revalidate_config(self) -> None:
+    def _revalidate_config(self, *, require_sample_weight: bool = False) -> None:
         """Validate + normalize EVERY config param from current attributes.
 
         Called at ``__init__`` and again at the start of ``fit()`` (the
@@ -411,9 +389,17 @@ def _revalidate_config(self) -> None:
         """
         self.anticipation = validate_anticipation(self.anticipation)
         _validate_learner_spec(
-            self.propensity_learner, kind="classifier", param_name="propensity_learner"
+            self.propensity_learner,
+            kind="classifier",
+            param_name="propensity_learner",
+            require_sample_weight=require_sample_weight,
+        )
+        _validate_learner_spec(
+            self.outcome_learner,
+            kind="regressor",
+            param_name="outcome_learner",
+            require_sample_weight=require_sample_weight,
         )
-        _validate_learner_spec(self.outcome_learner, kind="regressor", param_name="outcome_learner")
         # Specs stored VERBATIM (a passed learner object is the same object
         # in get_params()); fit-time make_learner does the resolution.
         self.n_folds = _validate_n_folds(self.n_folds)
@@ -475,13 +461,15 @@ def _validate_and_prepare(
         time: str,
         first_treat: str,
         covariates: Optional[Iterable[str]],
+        *,
+        require_sample_weight: bool = False,
     ) -> Tuple[pd.DataFrame, List[str]]:
         """Validate inputs; return the numeric working frame + covariate list."""
         # FULL config re-validation FIRST (mutation defense; anticipation
         # leads inside _revalidate_config — the ordering is load-bearing for
         # the anticipation-policy suite, which fits a bare DataFrame and
         # requires the config error to precede column checks).
-        self._revalidate_config()
+        self._revalidate_config(require_sample_weight=require_sample_weight)
 
         # covariates are REQUIRED (Chang's estimator exists for the
         # high-dimensional-X setting).
@@ -1267,7 +1255,7 @@ def _compute_dml_gt(
                     D_cell,
                     folds,
                     predict_method="predict_proba",
-                    context_label=f"{context} propensity",
+                    context_label=f"{context} propensity_learner",
                     sample_weight=w_cell,
                 )
                 or_res = cross_fit_predict(
@@ -1277,7 +1265,7 @@ def _compute_dml_gt(
                     folds,
                     predict_method="predict",
                     fit_mask=(D_cell == 0.0),
-                    context_label=f"{context} outcome",
+                    context_label=f"{context} outcome_learner",
                     sample_weight=w_cell,
                 )
         except DegenerateFoldError as exc:
@@ -1678,7 +1666,7 @@ def _compute_dml_rcs_gt(
                     D_cell,
                     folds,
                     predict_method="predict_proba",
-                    context_label=f"{context} propensity",
+                    context_label=f"{context} propensity_learner",
                     sample_weight=w_cell,
                 )
                 r_cell = (T_cell - lam_hat) * y_cell
@@ -1689,7 +1677,7 @@ def _compute_dml_rcs_gt(
                     folds,
                     predict_method="predict",
                     fit_mask=(D_cell == 0.0),
-                    context_label=f"{context} outcome",
+                    context_label=f"{context} outcome_learner",
                     sample_weight=w_cell,
                 )
         except DegenerateFoldError as exc:
@@ -1895,8 +1883,16 @@ def fit(
             i.i.d. sampling — Theorem 2's coverage claim does not carry
             over (REGISTRY DMLDiD Notes).
         """
+        self.results_ = None
+        self.is_fitted_ = False
         df, covariates = self._validate_and_prepare(
-            data, outcome, unit, time, first_treat, covariates
+            data,
+            outcome,
+            unit,
+            time,
+            first_treat,
+            covariates,
+            require_sample_weight=survey_design is not None,
         )
 
         # --- Survey/cluster resolution (CS transliteration, staggered.py) ---
@@ -2014,17 +2010,6 @@ def fit(
                 )
 
         weighted_moments = survey_design is not None
-        if weighted_moments:
-            # Learner capability gate: cross_fit_predict passes sample_weight
-            # BY KEYWORD, so a user learner without a keyword-addressable
-            # sample_weight (or **kwargs) would raise a raw TypeError mid-fit.
-            for spec, pname in (
-                (self.propensity_learner, "propensity_learner"),
-                (self.outcome_learner, "outcome_learner"),
-            ):
-                if isinstance(spec, str):
-                    continue  # native learners all accept sample_weight
-                _validate_learner_sample_weight_support(spec, pname)
 
         if self.panel:
             precomputed = self._precompute(

docs/api/dml_did.rst

@@ -141,6 +141,12 @@ estimator object plugs in directly; string names select library defaults
 only. :class:`~diff_diff.SieveLearner` is the exported configurable
 learner (``DMLDiD(outcome_learner=SieveLearner(k_max=3))``).
 
+Custom learner templates must support ``copy.deepcopy`` and return a distinct
+top-level object. DMLDiD checks this before fitting any cell and raises a
+targeted ``TypeError`` if copying fails or returns the original object. A
+custom ``__deepcopy__`` implementation remains responsible for isolating its
+nested mutable state.
+
 With ``seed`` set, fits are reproducible with the library's deterministic
 built-in learners; a user-supplied STOCHASTIC learner must additionally be
 seeded by the user (e.g. sklearn ``random_state``).
@@ -210,10 +216,11 @@ Restrictions
   base-period covariate). One consolidated ``UserWarning`` reports the
   drops.
 - **Degenerate cells skip loudly** — a cell that cannot be cross-fitted
-  (fewer members than folds, a singleton treated/control stratum, a
-  fail-closed learner error) is recorded as a NaN cell with a
-  machine-readable ``skip_reason`` and reported in a consolidated
-  warning; surviving cells still aggregate.
+  (fewer members than folds, a singleton treated/control stratum, or a
+  fold-time learner ``ValueError``) is recorded as a NaN cell with a
+  machine-readable ``skip_reason`` and reported in a consolidated warning;
+  surviving cells still aggregate. Learner-configuration errors, including
+  an uncloneable template, raise ``TypeError`` before any cell is estimated.
 - **Event-study surface is post-fit only** — fit-time
   ``event_study_effects`` is never populated; call
   ``results.aggregate('event_study')``.

docs/doc-deps.yaml

@@ -1457,7 +1457,7 @@ sources:
       - path: docs/methodology/REGISTRY.md
         section: "Cross-fitting, DR-score, and ridge infrastructure (DML)"
         type: methodology
-        note: "Duck-typed learner protocol (RegressorLearner/ClassifierLearner Protocols, validate_learner, _validate_predictions) + native learners (LinearLearner/RidgeLearner/LogitLearner/SieveLearner) wrapping linalg solvers. Contracts documented in REGISTRY: raw-X-no-intercept input, fit-reset semantics (documented limitation for stateful user learners), identified-columns prediction under rank deficiency."
+        note: "Duck-typed learner protocol (RegressorLearner/ClassifierLearner Protocols, validate_learner, _validate_predictions) + native learners (LinearLearner/RidgeLearner/LogitLearner/SieveLearner) wrapping linalg solvers. Contracts documented in REGISTRY: raw-X-no-intercept input, independent top-level deepcopy before cross-fitting, identified-columns prediction under rank deficiency."
       - path: docs/tutorials/32_dml_did.ipynb
         type: tutorial