From 9cd56498b68183381f61f44eab88e624b5614526 Mon Sep 17 00:00:00 2001 From: Charles Shaw Date: Thu, 3 Sep 2026 21:20:14 +0100 Subject: [PATCH 1/3] fix(dml): fail closed when cross-fit learners cannot be isolated Context: - Cross-fit reuse after a failed or identity-returning deepcopy can carry fitted state and prior-fold training data into a later fold. - DMLDiD must reject an invalid learner template before any group-time cell is fitted. Changes: - Require a distinct top-level deepcopy for cross-fit learner templates and raise a sanitised TypeError otherwise. - Probe custom DML learner specifications during configuration validation while retaining the direct cross_fit_predict backstop. - Document the top-level clone boundary and cover copy failures, identity copies, sklearn templates, exact OOF isolation, and pre-cell error propagation. Verification: - pytest -q tests/test_crossfit.py tests/test_dml_did.py tests/test_changelog_fragments.py tests/test_docs_ia.py tests/test_doc_deps_integrity.py - Python 3.9 targeted regression tests, including sklearn preflight - ruff, black, scoped mypy, changelog compiler, and native cross-fit output comparison against upstream/main - Not run: full default suite was stopped after an interrupted long-running attempt. --- ...20260903-dml-crossfit-learner-isolation.md | 6 ++ diff_diff/_crossfit.py | 78 ++++++++++--------- diff_diff/_learners.py | 8 +- diff_diff/dml_did.py | 8 +- docs/api/dml_did.rst | 15 +++- docs/methodology/REGISTRY.md | 19 +++-- tests/test_crossfit.py | 48 ++++++++++-- tests/test_dml_did.py | 54 +++++++++---- 8 files changed, 166 insertions(+), 70 deletions(-) create mode 100644 changelog.d/20260903-dml-crossfit-learner-isolation.md diff --git a/changelog.d/20260903-dml-crossfit-learner-isolation.md b/changelog.d/20260903-dml-crossfit-learner-isolation.md new file mode 100644 index 000000000..cdb5423f3 --- /dev/null +++ b/changelog.d/20260903-dml-crossfit-learner-isolation.md @@ -0,0 +1,6 @@ +### Fixed +- **Cross-fit learner isolation now fails closed**: DMLDiD rejects custom + learner templates whose `deepcopy` fails or returns the original object + before fitting any group-time cell, preventing reuse of the supplied + template across folds. Custom `__deepcopy__` implementations remain + responsible for isolating nested mutable state. diff --git a/diff_diff/_crossfit.py b/diff_diff/_crossfit.py index a2bca67d0..897f47590 100644 --- a/diff_diff/_crossfit.py +++ b/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): @@ -33,7 +34,6 @@ import copy import pickle -import warnings from dataclasses import dataclass, field from typing import Any, Dict, Iterator, Literal, Optional, Tuple, cast, overload @@ -58,38 +58,46 @@ _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, *, param_name: str) -> None: + """Fail validation unless a learner template deep-copies independently.""" + _clone_learner_template(learner, label=param_name) + + +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 +555,7 @@ 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) # 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 --git a/diff_diff/_learners.py b/diff_diff/_learners.py index 3a2fc495a..64ae95370 100644 --- a/diff_diff/_learners.py +++ b/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 --git a/diff_diff/dml_did.py b/diff_diff/dml_did.py index aa0e15a2f..05be2c878 100644 --- a/diff_diff/dml_did.py +++ b/diff_diff/dml_did.py @@ -41,7 +41,12 @@ 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, + assign_folds, + cross_fit_predict, +) from diff_diff._dr_scores import ( _chang_rcs_score_augmented_with_slope, chang_panel_score, @@ -114,6 +119,7 @@ def _validate_learner_spec(spec: Any, *, kind: str, param_name: str) -> None: ) return validate_learner(spec, kind=kind, param_name=param_name) + _probe_learner_cloneability(spec, param_name=param_name) def _validate_learner_sample_weight_support(spec: Any, param_name: str) -> None: diff --git a/docs/api/dml_did.rst b/docs/api/dml_did.rst index 38c0c61b4..863134028 100644 --- a/docs/api/dml_did.rst +++ b/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')``. diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 7e700e3c1..6edd391be 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -838,11 +838,14 @@ is estimator policy, applied by the consuming estimator). **Learner contract (`_learners.py`)** — duck-typed `RegressorLearner` / `ClassifierLearner` Protocols (sklearn-compatible `fit`/`predict`/ `predict_proba`); learners always receive raw `X` with NO intercept column and -manage the intercept internally; `fit` fully re-initializes (fit-reset -semantics). Cross-fitting fits a DEEP COPY of the user's never-fit learner -template in every fold, so no state — nested estimators and container -parameters included — can carry across folds; an un-deep-copyable learner is -reused with a loud `UserWarning` (the only residual reliance on fit-reset). Native +manage the intercept internally. Cross-fitting requires the user's never-fit +template to `deepcopy` to a distinct top-level object for every fold; an +uncopyable template, or one whose `__deepcopy__` returns itself, fails closed +with `TypeError` rather than being reused. +- **Note:** Custom `__deepcopy__` implementations are trusted to isolate nested + mutable state; the library checks only that the top-level copy is distinct. + +Native learners (`"linear"`, `"ridge"`, `"logit"`, `"sieve"`) wrap `solve_ols` / `solve_ridge` / `solve_logit` and the EfficientDiD polynomial sieve basis. Rank deficiency: the unpenalized fixed-design learners `LinearLearner` and @@ -2979,8 +2982,10 @@ the finite-dimensional `p_0` is handled by the variance correction below. meanings; `zero_weight_mass` (declared-survey fits only, the CS meaning) = a required group has rows but zero survey mass, so the weighted p̂/λ̂ would leave (0, 1); `cross_fit_degenerate` - = fold assignment or a fail-closed learner made the cell un-cross-fittable - (chained learner message quoted in the consolidated skip warning); + = fold assignment or a fold-time learner `ValueError` made the cell + un-cross-fittable (chained learner message quoted in the consolidated skip + warning); a learner-configuration error, including an uncloneable template, + raises `TypeError` before any cell exists; `non_finite_score` = the score/variance computation produced or received non-finite values. NaN cells carry NO influence-function payload entry. - **Note:** Covariates REQUIRED — `fit(covariates=None/[])` raises with a diff --git a/tests/test_crossfit.py b/tests/test_crossfit.py index f593b7549..b77286939 100644 --- a/tests/test_crossfit.py +++ b/tests/test_crossfit.py @@ -337,11 +337,13 @@ def predict(self, X): y = np.zeros(n) folds = assign_folds(n, 2, rng=_rng(20)) y[folds.test_mask(0)] = 1.0 # fold 0's outcomes differ from fold 1's - res = cross_fit_predict(AccumulatingMean(), X, y, folds) + template = AccumulatingMean() + res = cross_fit_predict(template, X, y, folds) # Each fold's prediction equals a FRESH fit on its complement alone. for k, train, test in folds.iter_folds(): expected = float(np.mean(y[train])) - np.testing.assert_allclose(res.oof_predictions[test], expected) + np.testing.assert_array_equal(res.oof_predictions[test], expected) + assert template._seen_y == [] def test_composite_learner_isolated_per_fold(self): # The per-fold deep copy isolates composites too: the template (and @@ -432,10 +434,34 @@ def test_ridge_overflow_becomes_degenerate_fold_error(self): sample_weight=np.full(n, 1e10), ) - def test_copy_failure_warns_loudly(self): + def test_copy_failure_raises_targeted_error(self): class Undeepcopyable: + def __init__(self): + self.fit_calls = 0 + + def __deepcopy__(self, memo): + raise TypeError("token=SECRET-DCOPY /home/user/private.csv") + + def fit(self, X, y, sample_weight=None): + self.fit_calls += 1 + self.mean_ = float(np.mean(y)) + return self + + def predict(self, X): + return np.full(len(X), self.mean_) + + X, y = _reg_setup() + folds = assign_folds(len(y), 2, rng=_rng(22)) + template = Undeepcopyable() + with pytest.raises(TypeError, match="Undeepcopyable.*TypeError") as exc_info: + cross_fit_predict(template, X, y, folds) + assert "SECRET-DCOPY" not in str(exc_info.value) + assert template.fit_calls == 0 + + def test_self_returning_deepcopy_raises_targeted_error(self): + class SelfCopying: def __deepcopy__(self, memo): - raise TypeError("cannot deep-copy this learner") + return self def fit(self, X, y, sample_weight=None): self.mean_ = float(np.mean(y)) @@ -446,9 +472,19 @@ def predict(self, X): X, y = _reg_setup() folds = assign_folds(len(y), 2, rng=_rng(22)) - with pytest.warns(UserWarning, match="could not deep-copy"): - res = cross_fit_predict(Undeepcopyable(), X, y, folds) + with pytest.raises(TypeError, match="SelfCopying.*original object"): + cross_fit_predict(SelfCopying(), X, y, folds) + + def test_sklearn_estimator_passes_distinct_clone_contract(self): + pytest.importorskip("sklearn") + from sklearn.linear_model import LinearRegression + + X, y = _reg_setup() + folds = assign_folds(len(y), 2, rng=_rng(22)) + template = LinearRegression() + res = cross_fit_predict(template, X, y, folds) assert np.isfinite(res.oof_predictions).all() + assert not hasattr(template, "coef_") def test_result_picklable(self): X, y = _reg_setup() diff --git a/tests/test_dml_did.py b/tests/test_dml_did.py index 6c54a63d3..d1b5666f6 100644 --- a/tests/test_dml_did.py +++ b/tests/test_dml_did.py @@ -121,6 +121,17 @@ def test_learner_spec_errors_name_the_param(self): with pytest.raises(TypeError, match="propensity_learner"): DMLDiD(propensity_learner=object()) + def test_sklearn_learners_pass_cloneability_preflight(self): + pytest.importorskip("sklearn") + from sklearn.linear_model import LinearRegression, LogisticRegression + + est = DMLDiD( + propensity_learner=LogisticRegression(), + outcome_learner=LinearRegression(), + ) + assert type(est.propensity_learner) is LogisticRegression + assert type(est.outcome_learner) is LinearRegression + def test_get_set_params_roundtrip_learner_object_identity(self): learner = SieveLearner(k_max=2) est = DMLDiD(outcome_learner=learner, seed=3) @@ -1080,6 +1091,28 @@ def test_mutated_config_raises_before_any_cell(self, data, attr, bad): with pytest.raises((ValueError, TypeError)): est.fit(data, **FIT_KW, **COV) + def test_mutated_uncopyable_learner_raises_before_any_cell(self, data): + class UndeepcopyableRegressor: + def __init__(self): + self.fit_calls = 0 + + def __deepcopy__(self, memo): + raise TypeError("cannot clone this learner") + + def fit(self, X, y): + self.fit_calls += 1 + return self + + def predict(self, X): + return np.zeros(len(X)) + + learner = UndeepcopyableRegressor() + est = DMLDiD(seed=0) + est.outcome_learner = learner + with pytest.raises(TypeError, match="outcome_learner.*UndeepcopyableRegressor.*TypeError"): + est.fit(data, **FIT_KW, **COV) + assert learner.fit_calls == 0 + class TestReportingWeightingLabel: def test_target_parameter_names_complete_case_weighting(self, fitted): @@ -1600,9 +1633,9 @@ def test_summary_uses_z_labels(self, fitted): assert "z-stat" in s and "P>|z|" in s assert "t-stat" not in s and "P>|t|" not in s - def test_deepcopy_failure_message_not_leaked(self, data): - # A foreign learner whose __deepcopy__ raises with sensitive text: - # the reuse warning names only the exception CLASS. + def test_deepcopy_failure_raises_sanitized_error(self, data): + # A foreign learner whose __deepcopy__ raises with sensitive text + # fails closed during configuration validation, before any cell fit. class LeakyDeepcopy: def __deepcopy__(self, memo): raise ValueError("token=SECRET-DCOPY /home/user/x.csv") @@ -1614,16 +1647,11 @@ def fit(self, X, y): def predict(self, X): return np.full(len(X), self.m) - with warnings.catch_warnings(record=True) as rec: - warnings.simplefilter("always") - res = DMLDiD(outcome_learner=LeakyDeepcopy(), seed=0).fit(data, **FIT_KW, **COV) - texts = [str(w.message) for w in rec] - assert not any("SECRET-DCOPY" in s for s in texts) - assert any("could not deep-copy" in s and "ValueError" in s for s in texts) - import json - - assert "SECRET-DCOPY" not in json.dumps(res.to_dict()) - assert "SECRET-DCOPY" not in res.summary() + with pytest.raises( + TypeError, match="outcome_learner.*LeakyDeepcopy.*ValueError" + ) as exc_info: + DMLDiD(outcome_learner=LeakyDeepcopy(), seed=0) + assert "SECRET-DCOPY" not in str(exc_info.value) def test_bootstrap_summary_labels_percentile_p(self, data): with warnings.catch_warnings(): From bb187a35ccb8c6df8c7244402d90fd9583360197 Mon Sep 17 00:00:00 2001 From: Charles Shaw Date: Thu, 3 Sep 2026 22:08:01 +0100 Subject: [PATCH 2/3] fix(dml): complete fail-closed clone isolation contract Context: - Cross-fit learner templates must fail closed when cloning cannot provide an independent, usable learner, without exposing stale DML results after a failed re-fit. Changes: - Validate clone protocol and weighted-fit capability during DML preflight and for each direct cross-fit fold; delayed errors name the public nuisance parameter. - Clear DMLDiD fitted state before every fit attempt, and cover panel, repeated-cross-section, survey, and direct helper failure paths. - Correct the DML tutorial and documentation registry, and record the warning-to-TypeError behavioural change. Verification: - /tmp/diff-diff-814-python39/bin/python -m pytest -q tests/test_crossfit.py tests/test_dml_did.py tests/test_survey_dml.py tests/test_changelog_fragments.py tests/test_docs_ia.py tests/test_doc_deps_integrity.py (665 passed, 6 skipped) - black --check, ruff check, and mypy --follow-imports=skip on touched source/test modules - python -m json.tool docs/tutorials/32_dml_did.ipynb and python .claude/scripts/changelog_compile.py check - Not run: full repository test suite. --- ...20260903-dml-crossfit-learner-isolation.md | 13 +-- diff_diff/_crossfit.py | 32 +++++- diff_diff/dml_did.py | 89 ++++++--------- docs/doc-deps.yaml | 2 +- docs/tutorials/32_dml_did.ipynb | 11 +- tests/test_crossfit.py | 66 +++++++++++- tests/test_dml_did.py | 101 +++++++++++++++++- 7 files changed, 242 insertions(+), 72 deletions(-) diff --git a/changelog.d/20260903-dml-crossfit-learner-isolation.md b/changelog.d/20260903-dml-crossfit-learner-isolation.md index cdb5423f3..3f848bb0a 100644 --- a/changelog.d/20260903-dml-crossfit-learner-isolation.md +++ b/changelog.d/20260903-dml-crossfit-learner-isolation.md @@ -1,6 +1,7 @@ -### Fixed -- **Cross-fit learner isolation now fails closed**: DMLDiD rejects custom - learner templates whose `deepcopy` fails or returns the original object - before fitting any group-time cell, preventing reuse of the supplied - template across folds. Custom `__deepcopy__` implementations remain - responsible for isolating nested mutable state. +### 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 --git a/diff_diff/_crossfit.py b/diff_diff/_crossfit.py index 897f47590..ff1bfb55e 100644 --- a/diff_diff/_crossfit.py +++ b/diff_diff/_crossfit.py @@ -33,6 +33,7 @@ """ import copy +import inspect import pickle from dataclasses import dataclass, field from typing import Any, Dict, Iterator, Literal, Optional, Tuple, cast, overload @@ -85,9 +86,31 @@ def _clone_learner_template(learner: Any, *, label: str) -> Any: ) -def _probe_learner_cloneability(learner: Any, *, param_name: str) -> None: - """Fail validation unless a learner template deep-copies independently.""" - _clone_learner_template(learner, label=param_name) +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: @@ -556,6 +579,9 @@ def cross_fit_predict( # (b) Learner errors during the fold -> DegenerateFoldError, chained. try: 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 --git a/diff_diff/dml_did.py b/diff_diff/dml_did.py index 05be2c878..c6f5faf4f 100644 --- a/diff_diff/dml_did.py +++ b/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 @@ -44,6 +43,7 @@ from diff_diff._crossfit import ( DegenerateFoldError, _probe_learner_cloneability, + _validate_sample_weight_support, assign_folds, cross_fit_predict, ) @@ -104,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 @@ -119,39 +121,9 @@ def _validate_learner_spec(spec: Any, *, kind: str, param_name: str) -> None: ) return validate_learner(spec, kind=kind, param_name=param_name) - _probe_learner_cloneability(spec, 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: @@ -404,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 @@ -417,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) @@ -481,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). @@ -1273,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( @@ -1283,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: @@ -1684,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 @@ -1695,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: @@ -1901,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) --- @@ -2020,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( diff --git a/docs/doc-deps.yaml b/docs/doc-deps.yaml index 7beb3e6c3..58b01f563 100644 --- a/docs/doc-deps.yaml +++ b/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 diff --git a/docs/tutorials/32_dml_did.ipynb b/docs/tutorials/32_dml_did.ipynb index 8c46a8193..6f5fc158e 100644 --- a/docs/tutorials/32_dml_did.ipynb +++ b/docs/tutorials/32_dml_did.ipynb @@ -491,11 +491,12 @@ " **both** nuisances converging at $o(N^{-1/4})$ - \"a fast learner cannot compensate\n", " a slow one\" - so with a deliberately misspecified propensity, the SEs and CIs in\n", " this table are illustrative rather than theory-backed.\n", - "- Custom learner objects are **deep-copied, never-fit, once per fold**. A learner\n", - " that cannot be deep-copied still runs, but with a loud `UserWarning`: the same\n", - " instance is REUSED across folds relying on its fit-reset behavior, so a\n", - " warm-start/stateful non-copyable learner can leak data across folds - make such\n", - " a learner fully re-initialize on every `fit()`. Under `survey_design=`\n", + "- A supplied custom learner template is **never fitted**. Cloneability is preflighted\n", + " before any group-time cell is fitted, and a fresh deep copy is fitted for each\n", + " fold of each estimable cell. A `deepcopy` that fails or returns the original\n", + " object raises `TypeError`; implement `__deepcopy__` to return an independent\n", + " instance. A custom implementation remains responsible for nested mutable state.\n", + " Under `survey_design=`\n", " (section 7) the `sample_weight` keyword becomes mandatory, and stochastic\n", " learners need their own internal seeding (`seed=` pins folds, not your\n", " learner's RNG).\n" diff --git a/tests/test_crossfit.py b/tests/test_crossfit.py index b77286939..7a8099ff3 100644 --- a/tests/test_crossfit.py +++ b/tests/test_crossfit.py @@ -1,6 +1,7 @@ """Tests for unit-level K-fold cross-fitting (PR-B0).""" import pickle +import traceback import numpy as np import pytest @@ -455,7 +456,12 @@ def predict(self, X): template = Undeepcopyable() with pytest.raises(TypeError, match="Undeepcopyable.*TypeError") as exc_info: cross_fit_predict(template, X, y, folds) - assert "SECRET-DCOPY" not in str(exc_info.value) + error = exc_info.value + assert "SECRET-DCOPY" not in str(error) + assert error.__cause__ is None + assert error.__context__ is None + formatted = "".join(traceback.format_exception(type(error), error, error.__traceback__)) + assert "SECRET-DCOPY" not in formatted assert template.fit_calls == 0 def test_self_returning_deepcopy_raises_targeted_error(self): @@ -475,6 +481,64 @@ def predict(self, X): with pytest.raises(TypeError, match="SelfCopying.*original object"): cross_fit_predict(SelfCopying(), X, y, folds) + def test_fold_clone_must_keep_the_learner_protocol(self): + class CloneLosesProtocol: + def __init__(self): + self.deepcopy_calls = 0 + self.fit_calls = 0 + + def __deepcopy__(self, memo): + self.deepcopy_calls += 1 + if self.deepcopy_calls == 2: + return object() + return type(self)() + + def fit(self, X, y, sample_weight=None): + self.fit_calls += 1 + self.mean_ = float(np.mean(y)) + return self + + def predict(self, X): + return np.full(len(X), self.mean_) + + X, y = _reg_setup() + folds = assign_folds(len(y), 2, rng=_rng(22)) + template = CloneLosesProtocol() + with pytest.raises(TypeError, match="fold 1 learner.*object.*fit"): + cross_fit_predict(template, X, y, folds) + assert template.deepcopy_calls == 2 + assert template.fit_calls == 0 + + def test_weighted_fold_clone_must_accept_sample_weight(self): + class CloneWithoutSampleWeight: + def fit(self, X, y): + self.mean_ = float(np.mean(y)) + return self + + def predict(self, X): + return np.full(len(X), self.mean_) + + class Template: + def __init__(self): + self.fit_calls = 0 + + def __deepcopy__(self, memo): + return CloneWithoutSampleWeight() + + def fit(self, X, y, sample_weight=None): + self.fit_calls += 1 + return self + + def predict(self, X): + return np.zeros(len(X)) + + X, y = _reg_setup() + folds = assign_folds(len(y), 2, rng=_rng(22)) + template = Template() + with pytest.raises(TypeError, match="fold 0 learner.*sample_weight"): + cross_fit_predict(template, X, y, folds, sample_weight=np.ones(len(y))) + assert template.fit_calls == 0 + def test_sklearn_estimator_passes_distinct_clone_contract(self): pytest.importorskip("sklearn") from sklearn.linear_model import LinearRegression diff --git a/tests/test_dml_did.py b/tests/test_dml_did.py index d1b5666f6..f04b2f068 100644 --- a/tests/test_dml_did.py +++ b/tests/test_dml_did.py @@ -5,6 +5,7 @@ """ import json +import traceback import warnings import numpy as np @@ -77,6 +78,20 @@ def test_n_folds_numpy_int_coerced(self): est = DMLDiD(n_folds=np.int64(5)) assert type(est.n_folds) is int + def test_clone_must_keep_the_required_learner_protocol(self): + class CloneWithoutPredict: + def __deepcopy__(self, memo): + return object() + + def fit(self, X, y, sample_weight=None): + return self + + def predict(self, X): + return np.zeros(len(X)) + + with pytest.raises(TypeError, match="outcome_learner.*object.*fit"): + DMLDiD(outcome_learner=CloneWithoutPredict()) + @pytest.mark.parametrize("bad", [0.0, 1.0, -0.1, 1.1, True, "0.05", np.nan]) def test_alpha_bounds_both_sides(self, bad): with pytest.raises(ValueError, match="alpha"): @@ -1113,6 +1128,85 @@ def predict(self, X): est.fit(data, **FIT_KW, **COV) assert learner.fit_calls == 0 + @pytest.mark.parametrize("panel,data_fixture", [(True, "data"), (False, "rcs_data")]) + def test_fold_copy_failure_propagates_as_a_hard_error(self, request, panel, data_fixture): + class CopyFailsDuringCrossFit: + def __init__(self): + self.deepcopy_calls = 0 + self.fit_calls = 0 + + def __deepcopy__(self, memo): + self.deepcopy_calls += 1 + if self.deepcopy_calls == 3: + raise OSError("token=SECRET-FOLD-COPY /home/user/private.csv") + return type(self)() + + def fit(self, X, y, sample_weight=None): + self.fit_calls += 1 + self.mean_ = float(np.mean(y)) + return self + + def predict(self, X): + return np.full(len(X), self.mean_) + + learner = CopyFailsDuringCrossFit() + data = request.getfixturevalue(data_fixture) + est = DMLDiD(n_folds=2, panel=panel, seed=0) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + previous = est.fit(data, **FIT_KW, **COV) + assert est.results_ is previous and est.is_fitted_ is True + est.outcome_learner = learner + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + with pytest.raises( + TypeError, + match="DMLDiD.*outcome_learner.*fold 1.*CopyFailsDuringCrossFit.*OSError", + ) as exc_info: + est.fit(data, **FIT_KW, **COV) + error = exc_info.value + assert error.__cause__ is None + assert error.__context__ is None + formatted = "".join(traceback.format_exception(type(error), error, error.__traceback__)) + assert "SECRET-FOLD-COPY" not in formatted + assert learner.fit_calls == 0 + assert learner.deepcopy_calls == 3 + assert est.results_ is None and est.is_fitted_ is False + assert not any("cross_fit_degenerate" in str(w.message) for w in recorded) + + def test_survey_clone_must_accept_sample_weight_before_any_cell(self, data): + from diff_diff.survey import SurveyDesign + + class CloneWithoutSampleWeight: + def fit(self, X, y): + self.mean_ = float(np.mean(y)) + return self + + def predict(self, X): + return np.full(len(X), self.mean_) + + class Template: + def __init__(self): + self.fit_calls = 0 + + def __deepcopy__(self, memo): + return CloneWithoutSampleWeight() + + def fit(self, X, y, sample_weight=None): + self.fit_calls += 1 + return self + + def predict(self, X): + return np.zeros(len(X)) + + learner = Template() + df = data.assign(w=1.0) + with pytest.raises(TypeError, match="outcome_learner.*sample_weight"): + DMLDiD(outcome_learner=learner, seed=0).fit( + df, **FIT_KW, **COV, survey_design=SurveyDesign(weights="w") + ) + assert learner.fit_calls == 0 + class TestReportingWeightingLabel: def test_target_parameter_names_complete_case_weighting(self, fitted): @@ -1651,7 +1745,12 @@ def predict(self, X): TypeError, match="outcome_learner.*LeakyDeepcopy.*ValueError" ) as exc_info: DMLDiD(outcome_learner=LeakyDeepcopy(), seed=0) - assert "SECRET-DCOPY" not in str(exc_info.value) + error = exc_info.value + assert "SECRET-DCOPY" not in str(error) + assert error.__cause__ is None + assert error.__context__ is None + formatted = "".join(traceback.format_exception(type(error), error, error.__traceback__)) + assert "SECRET-DCOPY" not in formatted def test_bootstrap_summary_labels_percentile_p(self, data): with warnings.catch_warnings(): From 50109fedb63fc3caf6c327c518b2877deb3e4ca0 Mon Sep 17 00:00:00 2001 From: igerber Date: Fri, 4 Sep 2026 07:27:11 -0400 Subject: [PATCH 3/3] docs(dml): narrow the clone-failure timing claim to the one-copy preflight The preflight guarantees a failing or self-returning __deepcopy__ raises before any cell is fitted; a per-fold clone failure later in the fit propagates as the same sanitized TypeError. State that on every surface that previously promised the stronger before-any-cell guarantee. --- .../20260903-dml-crossfit-learner-isolation.md | 12 +++++++----- docs/api/dml_did.rst | 10 +++++++--- docs/methodology/REGISTRY.md | 4 +++- docs/tutorials/32_dml_did.ipynb | 3 ++- 4 files changed, 19 insertions(+), 10 deletions(-) diff --git a/changelog.d/20260903-dml-crossfit-learner-isolation.md b/changelog.d/20260903-dml-crossfit-learner-isolation.md index 3f848bb0a..a52cfbe7d 100644 --- a/changelog.d/20260903-dml-crossfit-learner-isolation.md +++ b/changelog.d/20260903-dml-crossfit-learner-isolation.md @@ -1,7 +1,9 @@ ### 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. + warning-and-reuse fallback. A one-copy preflight raises `TypeError` before + any group-time cell is fitted when a custom learner template's `deepcopy` + fails or returns the original object; a per-fold clone failure later in the + fit propagates as the same sanitized `TypeError` (never a NaN-cell skip). + 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 --git a/docs/api/dml_did.rst b/docs/api/dml_did.rst index 863134028..8d0f8c252 100644 --- a/docs/api/dml_did.rst +++ b/docs/api/dml_did.rst @@ -142,8 +142,10 @@ 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 +top-level object. DMLDiD preflights one copy of each template before fitting +any cell and raises a targeted ``TypeError`` if copying fails or returns the +original object; a clone failure in a later fold propagates as the same +``TypeError`` rather than a NaN-cell skip. A custom ``__deepcopy__`` implementation remains responsible for isolating its nested mutable state. @@ -220,7 +222,9 @@ Restrictions 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. + an uncloneable template, raise ``TypeError`` at the preflight before any + cell is estimated; a per-fold clone failure later in the fit propagates as + a hard ``TypeError`` too, never as a NaN cell. - **Event-study surface is post-fit only** — fit-time ``event_study_effects`` is never populated; call ``results.aggregate('event_study')``. diff --git a/docs/methodology/REGISTRY.md b/docs/methodology/REGISTRY.md index 6edd391be..0ee24bf13 100644 --- a/docs/methodology/REGISTRY.md +++ b/docs/methodology/REGISTRY.md @@ -2985,7 +2985,9 @@ the finite-dimensional `p_0` is handled by the variance correction below. = fold assignment or a fold-time learner `ValueError` made the cell un-cross-fittable (chained learner message quoted in the consolidated skip warning); a learner-configuration error, including an uncloneable template, - raises `TypeError` before any cell exists; + raises `TypeError` at the one-copy preflight before any cell exists, and a + per-fold clone failure later in the fit propagates as a hard `TypeError` + (never a NaN cell); `non_finite_score` = the score/variance computation produced or received non-finite values. NaN cells carry NO influence-function payload entry. - **Note:** Covariates REQUIRED — `fit(covariates=None/[])` raises with a diff --git a/docs/tutorials/32_dml_did.ipynb b/docs/tutorials/32_dml_did.ipynb index 6f5fc158e..8b968a0b9 100644 --- a/docs/tutorials/32_dml_did.ipynb +++ b/docs/tutorials/32_dml_did.ipynb @@ -494,7 +494,8 @@ "- A supplied custom learner template is **never fitted**. Cloneability is preflighted\n", " before any group-time cell is fitted, and a fresh deep copy is fitted for each\n", " fold of each estimable cell. A `deepcopy` that fails or returns the original\n", - " object raises `TypeError`; implement `__deepcopy__` to return an independent\n", + " object raises `TypeError` (at the preflight, or as a hard error from a later\n", + " fold - never a silent NaN cell); implement `__deepcopy__` to return an independent\n", " instance. A custom implementation remains responsible for nested mutable state.\n", " Under `survey_design=`\n", " (section 7) the `sample_weight` keyword becomes mandatory, and stochastic\n",