DML cross-fitting can reuse stateful custom learners across folds

[shawcharles]

Problem

DML cross-fitting currently reuses a custom learner template when copy.deepcopy() fails. That breaks the out-of-fold isolation contract for a stateful or warm-start learner: state from an earlier training complement, including observations that are evaluation data in another fold, can affect later predictions.

There is a second silent variant: a custom __deepcopy__ can return self. The current code accepts that object as a fold learner without warning.

Reproduction

At current main (e11d414), I used an eight-observation regression problem with two deterministic folds and a learner that accumulates all outcomes passed to fit().

A fresh learner per fold should produce:

fold ids:       [0, 1, 0, 0, 0, 1, 1, 1]
fresh OOF:      [4.75, 2.25, 4.75, 4.75, 4.75, 2.25, 2.25, 2.25]

An uncopyable learner currently produces:

actual OOF:     [4.75, 3.50, 4.75, 4.75, 4.75, 3.50, 3.50, 3.50]
max difference: 1.25
template seen:  8 outcomes

It emits a warning saying it is reusing the instance. A learner whose __deepcopy__ returns self produces the same contaminated OOF predictions and template mutation, but emits no warning.

The direct path is diff_diff/_crossfit.py::_fresh_learner (lines 61-92), which catches a copy failure and returns the supplied template. cross_fit_predict() then calls it immediately before each fold fit (around lines 520-585). DMLDiD calls this path for both propensity and outcome nuisance learners. The public documentation currently accepts any duck-typed learner with fit plus predict or predict_proba (docs/api/dml_did.rst, learner configuration).

The existing tests/test_crossfit.py::test_copy_failure_warns_loudly codifies the warning-and-reuse behaviour.

Proposed boundary

Make learner isolation fail closed:

1. A cross-fit learner template must deep-copy to a distinct top-level object for every fold.
2. If copying raises, or returns the original object, raise a targeted TypeError rather than fitting the template.
3. Preflight all fold clones before fitting any clone, so a later copy failure cannot leave an earlier fold fitted. This retains at most n_folds unfitted clones temporarily.
4. Include the caller/nuisance context and fold index, learner class, and copy exception class in the error. Do not expose the foreign exception message.
5. Capture only the foreign exception class inside the except block, then raise after leaving it. raise ... from None is insufficient for the sanitisation boundary because the original exception remains reachable through __context__.
6. Document that custom __deepcopy__ implementations are trusted to isolate nested mutable state. Top-level identity checking cannot prove arbitrary nested aliases without a new clone/factory protocol.

Acceptance criteria

• An uncopyable learner and a self-returning __deepcopy__ fail before any learner fit() call.
• A second-copy failure proves the clone preflight is transactional.
• Copyable stateful learners match an independent fresh-instance OOF oracle.
• Panel and repeated-cross-section DML paths cover propensity and outcome learner failures.
• The sanitised exception contains no foreign message in args, __context__, __cause__, or formatted traceback.
• Existing built-in learner outputs remain unchanged.
• The learner API documentation and methodology registry describe the cloneability requirement, with a changelog fragment.

This changes the current warning-and-reuse fallback, so I would like to confirm the boundary before preparing a narrowly scoped PR.

[igerber]

Thanks for the careful write-up and reproduction, and for your continued support of the library. I confirmed both variants locally on e11d414: an accumulating learner whose __deepcopy__ raises gets reused with the warning, and one whose __deepcopy__ returns self gets reused with no warning at all. In both cases the second fold's predictions carry the first fold's training data and the user's template is mutated. The silent identity-returning case is the more serious of the two, and it conflicts with the library's no-silent-failures rule. The warn-and-reuse fallback was a deliberate choice in #789, but you have made a good case that it is the wrong trade for a DML estimator, so I agree with failing closed.

Confirming the boundary, with a few adjustments:

Agreed as proposed

• A cross-fit learner template must deep-copy to a distinct top-level object. Copy failure, or __deepcopy__ returning the original object, raises a targeted TypeError. No opt-out knob; the remedy is a __deepcopy__ that yields an independent instance, and the error message should say so.
• The error names the nuisance parameter (propensity_learner / outcome_learner), the learner class, and the copy exception class only. Capturing the class inside the except and raising outside it is fine.
• Custom __deepcopy__ implementations are trusted to isolate nested mutable state. Top-level identity is the boundary we can check; document it as such.
• Copyable stateful learners must match an independent fresh-instance out-of-fold oracle, asserting the exact fold values, not just finiteness.

Adjustments

1. Preflight site. Please put the probe in _validate_learner_spec (dml_did.py, alongside the existing duck-type check), not inside the fold loop. DMLDiD calls cross_fit_predict twice per (g, t) cell, so a per-fold check would only fire after earlier cells had already been fitted. One probe per learner at validation time (copy, check is not the template, discard) fails the fit before any cell, fold, or fit() call, and covers the panel and repeated-cross-section paths in one place. _fresh_learner should still raise rather than warn, as a backstop for direct callers of cross_fit_predict.
2. Drop the all-folds clone preflight and the "second-copy failure" test. With the up-front probe, per-fold copying can stay lazy. Holding n_folds unfitted clones at once is a real memory cost for large templates, and it buys nothing: the template is never fitted, so a later copy failure leaves no corrupt state.
3. Exception propagation. The per-cell handler in DMLDiD catches only DegenerateFoldError, so the new TypeError should propagate as a hard failure rather than becoming a NaN cell. That is the intended behavior; this is a configuration error, not fold-data degeneracy. Please pin that with a test.
4. Existing test. test_copy_failure_warns_loudly in tests/test_crossfit.py flips to pytest.raises(TypeError), with a sibling for the self-returning case. Please also add a test that a plain scikit-learn estimator passes the probe, since that is the dominant real-world custom learner (gate it with importorskip).

Surfaces to update together

• diff_diff/_crossfit.py module docstring and _fresh_learner docstring
• diff_diff/_learners.py header (the "accepted limitation" sentence)
• docs/methodology/REGISTRY.md, Learner contract paragraph (the "only residual reliance on fit-reset" sentence), using a - **Note:** line for the cloneability requirement
• docs/api/dml_did.rst, Learner configuration section
• a changelog.d/YYYYMMDD-<slug>.md fragment

Built-in learner outputs must remain bit-identical. With that scope, please go ahead and open the PR.