Skip to content

Commit 6874b5d

Browse files
authored
fix(dml): fail closed when cross-fit learners cannot be isolated (#816)
1 parent abcd785 commit 6874b5d

18 files changed

Lines changed: 534 additions & 232 deletions

‎.github/workflows/rust-test.yml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -154,7 +154,7 @@ jobs:
154154

155155
- name: Install test dependencies
156156
# Keep in sync with pyproject.toml [project.dependencies] and [project.optional-dependencies.dev]
157-
run: pip install pytest pytest-xdist numpy pandas scipy
157+
run: pip install pytest pytest-xdist numpy pandas scipy scikit-learn
158158

159159
- name: Install OpenBLAS (Linux)
160160
if: runner.os == 'Linux'
@@ -255,7 +255,7 @@ jobs:
255255

256256
- name: Install dependencies
257257
# Keep in sync with pyproject.toml [project.dependencies] and [project.optional-dependencies.dev]
258-
run: pip install numpy pandas scipy pytest pytest-xdist
258+
run: pip install numpy pandas scipy pytest pytest-xdist scikit-learn
259259

260260
- name: Verify pure Python mode
261261
run: |

‎DEFERRED.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,7 @@ decisions (refactor waivers, perf trade-offs, test-infrastructure calls) are rec
149149
| **MMM launch carousel carries scoped marketing claims, not exporter-contract documentation.** The deck (a LinkedIn marketing artifact, not a docs surface) states capabilities in scoped-but-punchy form: guardrail copy says "the easy mistakes fail loudly" / "you own the design, it owns the math" (the caller-owned estimand/population/window/outcome-scale alignment lives in REGISTRY.md's MMM section and the exporter docstrings, not on slides), and slide 8 shows tutorial 29's worked staggered-boost lift row WITHOUT an on-slide linearity qualifier - the compression's linear-channel scoping is the tutorial's job, which the CTA points to. Absolutes ("no silent mis-calibration", "any estimate exports", "anything with an estimate + SE") are ban-listed by `tests/test_mmm_carousel_claims.py`; local review R1/R2 pressed for compliance-style scoping language on the slides and the marketing-appropriate scoped copy was chosen instead. | `carousel/generate_mmm_carousel.py`, `tests/test_mmm_carousel_claims.py` | mmm-carousel / 2026-08-20 |
150150
| **MultiPeriodDiD deprecation shim loses static constructor-arg checking (3.9 window).** The M-010 shim is `__init__(*args, **kwargs)` + an import-time `__signature__` mirror of DiD's constructor: runtime introspection (get_params/set_params, `inspect.signature`) and eager validation are fully preserved, but static type checkers / IDEs cannot check constructor arguments for the deprecated class until the ledger-tracked M-010 removal. Accepted: the alternative (hand-mirroring ~20 parameters) is a drift magnet on a class with one minor version of remaining life. | `diff_diff/estimators.py` | 3(a) / 2026-08-07 |
151151
| **DCDH `sklearn.base.clone` param-identity failure won't-fix.** `ChaisemartinDHaultfoeuille._validate_paths_of_interest` unconditionally canonicalizes `paths_of_interest` into a fresh `List[Tuple[int, ...]]`, so sklearn `clone()`'s post-construction `param1 is param2` identity check fails for configured instances - a pre-existing normalization the BaseEstimator mixin PR documented rather than changed (get_params/set_params signatures are clone-compatible; the dependency-free `cls(**est.get_params())` config-equality contract is the enforced one, `tests/test_base_estimator.py`). Fixing would mean returning the caller's raw object from a validator whose job is canonicalization. | `chaisemartin_dhaultfoeuille.py` | mixin PR / 2026-08-01 |
152-
| **scikit-learn stays out of dev deps; clone-identity tests remain importorskip-only.** The sklearn-`clone()` round-trip tests (`test_base_estimator.py`, had/rdd/cic suites) run only where scikit-learn happens to be installed - deliberate, matching the numpy/pandas/scipy-only dependency posture; the always-running contract is the dependency-free re-instantiation config-equality test. | `tests/test_base_estimator.py` | mixin PR / 2026-08-01 |
152+
| **scikit-learn is a required dev/CI dependency, optional at runtime; clone-identity tests stay `importorskip`-guarded.** Supersedes the 2026-08-01 decision that kept scikit-learn out of the dev extra: #816 added it (unpinned) to the `dev` extra and to both `rust-test.yml` install steps, so the sklearn `clone()` round-trip tests (`test_base_estimator.py`, had/rdd/cic suites) and the DML cross-fit learner tests always run in CI. The library's runtime dependency posture is unchanged (numpy/pandas/scipy only) and the `importorskip` guards remain, so the suite still passes in an environment without scikit-learn; the always-running contract is still the dependency-free re-instantiation config-equality test. | `pyproject.toml`, `tests/test_base_estimator.py` | #816 |
153153
| **Plan-review hash gate threat model: accident prevention, NOT adversarial defense.** The ExitPlanMode content-hash gate (hook + `plan_snapshot.py`) exists to stop accidents — stale approvals, concurrent-worktree cross-talk, plans edited mid-review — all of which it closes by construction (snapshot identity + invocation-unique state tokens, 30+ behavioral tests). It does NOT and cannot defend against a malicious local process: nothing verifies review AUTHORSHIP, and such an actor can simply write a matching review file directly — no userland hook can prevent that short of signed reviews, which is out of scope for a personal workflow aid. Review findings that presuppose a hostile local actor against this gate are waived by this decision (2026-07-20, after 7 local AI-review rounds converged on ever-deeper adversarial-model refinements with no reachable fixpoint). Genuine accident vectors remain in scope and are fixed as found. | `.claude/hooks/check-plan-review.py`, `.claude/scripts/plan_snapshot.py` | 2026-07-20 |
154154
| **`StackedDiD` survey re-resolution intra-file dedup not warranted.** The cross-estimator ContinuousDiD/EfficientDiD panel-to-unit collapse consolidation landed in #226 (`ResolvedSurveyDesign.subset_to_units_by_row_idx` / `build_unit_first_row_index`); StackedDiD deliberately re-resolves at stacked granularity (control units are duplicated across sub-experiments). The residual intra-file dedup (raw-weight extraction ×3, compose-normalize ×3, resolve-on-stacked ×2) is stacked-specific, low value, and touches the numerically-sensitive composed-weight renormalization; cross-estimator unification was assessed and is not warranted (genuinely different mechanisms; all three already delegate to `_resolve_survey_for_fit` / `compute_survey_metadata`). | `stacked_did.py` | #226 review |
155155
| **`HeterogeneousAdoptionDiD` Phase-4 Pierce-Schott (2016) replication harness waived.** R parity at `atol=1e-8` on the same 3 DGPs is a strictly stronger anchor than reproducing Fig 2's pointwise CIs on the LBD-restricted PNTR panel (paper §5.2 self-acknowledges NP estimators too noisy there). Re-open only on user demand. See REGISTRY HAD Deviations Notes #3/#4. | `benchmarks/`, `tests/` | Phase 2a / 2026-05-20 |
Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
### Behavioral Changes
2+
- **Cross-fit learner isolation now fails closed**: DMLDiD replaces the prior
3+
warning-and-reuse fallback. A one-copy preflight raises `TypeError` before
4+
any group-time cell is fitted when a custom learner template's `deepcopy`
5+
fails or returns the original object; a per-fold clone failure later in the
6+
fit propagates as the same sanitized `TypeError` (never a NaN-cell skip).
7+
Implement `__deepcopy__` to return an independent instance; custom
8+
implementations remain responsible for nested mutable state. A failed
9+
re-fit now also clears any previous fitted result.
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
### Fixed
2+
- **Preserve omitted-parameter sentinels when copying**: shallow and deep
3+
copies of the shared deprecation sentinel retain its identity, keeping
4+
StackedDiD parameter dictionaries equal after scikit-learn cloning.
5+
Deprecated-parameter warnings and alias resolution are unchanged.

‎diff_diff/_crossfit.py‎

Lines changed: 63 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -10,11 +10,12 @@
1010
1111
``cross_fit_predict`` produces out-of-fold nuisance predictions for EVERY
1212
unit: for each fold k the learner is fit on ``train_mask(k) & fit_mask`` and
13-
predicts all units in fold k. Each fold fits a DEEP COPY of the user's
14-
(never-fit) learner template, so no state — nested estimators and container
15-
parameters included — can carry across folds; an un-deep-copyable learner is
16-
reused with a loud warning under the fit-reset contract (see
17-
``diff_diff._learners``).
13+
predicts all units in fold k. Each fold uses a distinct top-level DEEP COPY of
14+
the user's (never-fit) learner template, so the template itself cannot carry
15+
fitted state across folds. Custom ``__deepcopy__`` implementations remain
16+
responsible for isolating nested mutable state. A template that cannot be
17+
deep-copied to a distinct top-level object fails closed with ``TypeError``
18+
(see ``diff_diff._learners``).
1819
1920
Exception semantics (determinate):
2021
@@ -32,8 +33,8 @@
3233
"""
3334

3435
import copy
36+
import inspect
3537
import pickle
36-
import warnings
3738
from dataclasses import dataclass, field
3839
from typing import Any, Dict, Iterator, Literal, NamedTuple, Optional, Tuple, cast, overload
3940

@@ -60,41 +61,58 @@
6061
_LOG_LOSS_CLIP = 1e-15
6162

6263

63-
def _fresh_learner(learner: Any, *, stacklevel: int) -> Any:
64-
"""Per-fold learner isolation: a deep copy of the (never-fit) template.
65-
66-
``copy.deepcopy`` of the user's template gives every fold a fully
67-
independent learner — nested estimators, estimators inside lists/dicts,
68-
accumulators, and warm-start state included — so no state (and therefore
69-
no data from a previous complement, which includes the current evaluation
70-
fold) can carry across folds. This is strictly stronger than
71-
get_params-based reconstruction (which shares any estimator stored inside
72-
a container parameter). The template itself is never fit. A copy FAILURE
73-
is never silent: the instance is reused with a loud ``UserWarning`` naming
74-
the learner and the fit-reset assumption now being relied on
75-
(no-silent-failures rule). ``stacklevel`` is supplied by the caller so the
76-
warning is attributed to the entry point's own caller (``cross_fit_predict``
77-
passes the level that lands on the user's call site; internal consumers
78-
of ``iter_fold_fits`` pass the level that lands on themselves).
64+
def _clone_learner_template(learner: Any, *, label: str) -> Any:
65+
"""Return a distinct deep copy or raise a sanitised ``TypeError``.
66+
67+
The check deliberately proves only top-level identity. A custom
68+
``__deepcopy__`` remains responsible for isolating nested mutable state.
69+
Error messages expose the learner class and copy exception class, never
70+
foreign exception text that could carry credentials, paths, or data.
7971
"""
8072
try:
81-
return copy.deepcopy(learner)
82-
except Exception as exc: # noqa: BLE001 - loud fallback, never silent
83-
# Exception CLASS only, never the message: a foreign learner's
84-
# __deepcopy__ error text can embed credentials/paths/data excerpts,
85-
# and this warning lands in notebook/CI logs (the same boundary as
86-
# DMLDiD's persisted-diagnostics sanitization).
87-
warnings.warn(
88-
f"_crossfit: could not deep-copy the "
89-
f"{type(learner).__name__} template for this fold "
90-
f"({type(exc).__name__}); "
91-
"REUSING the same instance and relying on its fit-reset behavior. "
92-
"A warm-start/stateful learner in this situation can leak data "
93-
"across folds.",
94-
UserWarning,
95-
stacklevel=stacklevel,
73+
clone = copy.deepcopy(learner)
74+
except Exception as exc: # noqa: BLE001 - sanitize a foreign exception boundary
75+
copy_error_class = type(exc).__name__
76+
else:
77+
if clone is not learner:
78+
return clone
79+
raise TypeError(
80+
f"{label}: {type(learner).__name__} learner template's __deepcopy__ "
81+
"returned the original object; implement __deepcopy__ to return an "
82+
"independent instance."
9683
)
97-
return learner
84+
85+
raise TypeError(
86+
f"{label}: could not deep-copy {type(learner).__name__} learner template "
87+
f"({copy_error_class}); implement __deepcopy__ to return an independent instance."
88+
)
89+
90+
91+
def _probe_learner_cloneability(learner: Any, *, kind: str, param_name: str) -> Any:
92+
"""Return an independent clone after validating its learner protocol."""
93+
clone = _clone_learner_template(learner, label=param_name)
94+
validate_learner(clone, kind=kind, param_name=param_name)
95+
return clone
96+
97+
98+
def _validate_sample_weight_support(learner: Any, *, param_name: str) -> None:
99+
"""Require a learner ``fit`` method that accepts ``sample_weight`` by keyword."""
100+
try:
101+
sig = inspect.signature(learner.fit)
102+
except (TypeError, ValueError): # pragma: no cover - exotic callables
103+
return # Cannot introspect; let the learner surface any fit-time error.
104+
for param in sig.parameters.values():
105+
if param.kind is inspect.Parameter.VAR_KEYWORD:
106+
return
107+
if param.name == "sample_weight" and param.kind in (
108+
inspect.Parameter.POSITIONAL_OR_KEYWORD,
109+
inspect.Parameter.KEYWORD_ONLY,
110+
):
111+
return
112+
raise TypeError(
113+
f"{param_name}: learner object {type(learner).__name__!r} must accept "
114+
"sample_weight by keyword in fit(); add a sample_weight parameter (or **kwargs)."
115+
)
98116

99117

100118
def _unique_or_raise(arr: np.ndarray, name: str, **kwargs: Any) -> Any:
@@ -531,7 +549,6 @@ def _fit_subset(
531549
sample_weight: Optional[np.ndarray],
532550
k: int,
533551
label: str,
534-
warn_stacklevel: int,
535552
) -> Tuple[Any, int, Optional[np.ndarray]]:
536553
"""Fit a fresh deep copy of ``learner`` on ``fit_idx`` (fold ``k``).
537554
@@ -567,7 +584,10 @@ def _fit_subset(
567584

568585
# (b) Learner errors during the fold -> DegenerateFoldError, chained.
569586
try:
570-
fold_learner = _fresh_learner(learner, stacklevel=warn_stacklevel)
587+
fold_learner = _clone_learner_template(learner, label=f"{label}fold {k} learner")
588+
validate_learner(fold_learner, kind=kind, param_name=f"{label}fold {k} learner")
589+
if w_fit is not None:
590+
_validate_sample_weight_support(fold_learner, param_name=f"{label}fold {k} learner")
571591
# Unweighted path calls fit(X, y) WITHOUT the keyword: the
572592
# advertised duck-typed contract is fit/predict(_proba), so a
573593
# learner whose fit signature is only (X, y) must work when no
@@ -647,7 +667,6 @@ def _iter_fold_fits_impl(
647667
sample_weight: Optional[np.ndarray],
648668
kind: str,
649669
label: str,
650-
warn_stacklevel: int,
651670
) -> Iterator[FoldFit]:
652671
for k, train_idx, test_idx in folds.iter_folds():
653672
fit_idx = train_idx[fit_mask_arr[train_idx]]
@@ -661,7 +680,6 @@ def _iter_fold_fits_impl(
661680
sample_weight=sample_weight,
662681
k=k,
663682
label=label,
664-
warn_stacklevel=warn_stacklevel,
665683
)
666684
yield FoldFit(
667685
k=k,
@@ -689,7 +707,6 @@ def iter_fold_fits(
689707
predict_method: str = "predict",
690708
sample_weight: Optional[np.ndarray] = None,
691709
context_label: str = "",
692-
warn_stacklevel: int = 4,
693710
) -> Iterator[FoldFit]:
694711
"""Per-fold fitted learners (the building block under ``cross_fit_predict``).
695712
@@ -698,9 +715,8 @@ def iter_fold_fits(
698715
``k``, each holding a fresh deep copy of ``learner`` fit on
699716
``train_mask(k) & fit_mask``. Predictions are the caller's job via
700717
``_predict_subset``; this is what lets a consumer fit a NESTED nuisance on
701-
fold ``k``'s own training units. ``warn_stacklevel`` is the frame the
702-
deep-copy-failure warning is attributed to (default: the frame advancing
703-
the iterator).
718+
fold ``k``'s own training units. Clone failures raise a sanitised
719+
``TypeError`` when the iterator advances; the template is never reused.
704720
"""
705721
X64, y64, fit_mask_arr, sw, kind = _prepare_cross_fit_inputs(
706722
learner,
@@ -712,9 +728,7 @@ def iter_fold_fits(
712728
sample_weight=sample_weight,
713729
)
714730
label = f"{context_label}: " if context_label else ""
715-
return _iter_fold_fits_impl(
716-
learner, X64, y64, folds, fit_mask_arr, sw, kind, label, warn_stacklevel
717-
)
731+
return _iter_fold_fits_impl(learner, X64, y64, folds, fit_mask_arr, sw, kind, label)
718732

719733

720734
@overload
@@ -768,8 +782,6 @@ def cross_fit_predict(
768782
fold_losses = np.empty(folds.n_folds, dtype=np.float64)
769783
n_fit_per_fold = np.empty(folds.n_folds, dtype=np.int64)
770784

771-
# stacklevel 5: _fresh_learner -> _fit_subset -> generator frame ->
772-
# cross_fit_predict -> the user's call site.
773785
for ff in iter_fold_fits(
774786
learner,
775787
X,
@@ -779,7 +791,6 @@ def cross_fit_predict(
779791
predict_method=predict_method,
780792
sample_weight=sample_weight,
781793
context_label=context_label,
782-
warn_stacklevel=5,
783794
):
784795
n_fit_per_fold[ff.k] = ff.n_fit
785796
pred = _predict_subset(

‎diff_diff/_deprecation.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,13 @@ class _NotSupplied:
5252
def __repr__(self) -> str: # pragma: no cover - debugging aid
5353
return "<not supplied>"
5454

55+
# Parameter cloning must preserve this stateless marker's identity.
56+
def __copy__(self) -> "_NotSupplied":
57+
return self
58+
59+
def __deepcopy__(self, memo: Any) -> "_NotSupplied":
60+
return self
61+
5562

5663
NOT_SUPPLIED = _NotSupplied()
5764

‎diff_diff/_learners.py‎

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

0 commit comments

Comments
 (0)