[kbattocchi]

Summary

Fixes #1066.

Fix two errors in CausalForestDML's fit-time doubly robust training-data ATE/ATT inference:

• Count non-missing predictions along the sample axis (axis=0), separately for each outcome-treatment pair and after ATT masking. Previously, pooling counts across outputs or treatment contrasts understated standard errors.
• Preserve outcome-treatment ordering when constructing the doubly robust correction. The previous cross-product/reshape combination mixed the axes when both dimensions exceeded one, affecting point estimates as well as standard errors. Use an explicit einsum with sample, outcome, and treatment axes.

Regression coverage

• Deterministic aggregation checks cover vector, single-column, and multi-output outcomes; binary and three-arm treatments; missing rows and unequal column missingness; and masked subsets.
• Fitted-estimator checks cover ATE/ATT estimates, standard errors, confidence intervals, z-statistics, and p-values for every treatment group. They assert finite predictions and use ordinary means, standard deviations, and sample counts, with a comment explaining why the data and estimator seeds must be retained.
• Duplicate-outcome invariance checks compare public ATE/ATT point estimates, standard errors, and confidence intervals before and after duplicating an outcome, for two, three, and four treatment arms. They do not recalculate a reference standard-error formula.

The invariance test exposed the correction-order error. The existing known-ground-truth test uses duplicate outcomes, identical true effects for both non-control treatments, and absolute tolerances that allowed this error to escape; recomputing inference from the estimator's internal corrected predictions also cannot detect incorrect corrections upstream.

Validation

• All 45 focused regression cases plus the existing causal-forest ATE inference test passed (46 total).
• Independently restoring the original missing-axis bug in memory makes all three invariance cases fail.
• Independently restoring the original correction ordering in memory makes the three- and four-arm invariance cases fail; the binary-treatment control still passes.

[Copilot]

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Copilot review overview

3 open issues

• With the new per-column nonnan, columns that are entirely NaN will have nonnan == 0, which can… · New
• These tests require OOB predictions to be finite for every observation (`assert… · New
• For clarity and to lock the intended convention to match np.nanstd(..., ddof=0) used in the… · New

What changed in this PR

Fixes fit-time doubly robust ATE/ATT inference in CausalForestDML for multi-output outcomes / multi-arm treatments by correcting both the per-output sample counting for standard errors and the axis ordering in the DR correction term.

Changes:

• Count non-missing DR predictions per outcome-treatment column (and after ATT masking) when computing standard errors.
• Replace cross_product(...).reshape(...) with an explicit einsum to preserve (sample, outcome, treatment) axis ordering in the DR correction.
• Add regression tests covering multi-output, multi-arm, missingness patterns, ATT masking, and duplicate-outcome invariance.

File	Description
econml/​dml/​causal_forest.py	Fixes DR correction axis mixing and computes per-column non-missing counts for ATE/ATT stderr.
econml/​tests/​test_dml.py	Adds targeted regressions validating ATE/ATT point estimates, stderr, and inference outputs across shapes/arms/missingness.

_{Review effort: Lite}

[kbattocchi]

I checked both versions on an all-NaN column: both return NaN and emit the same nanmean/nanstd warnings; the new denominator does not introduce an additional division warning. Replacing zero denominators would not suppress those reduction warnings. I'd retain the diagnostic rather than silently hide an unavailable estimate. I've added regression cases for all-missing columns, with and without a sample mask, checking NaN results and unaffected valid columns. That test-only follow-up is committed locally, pending publication approval.

[kbattocchi]

Published the regression coverage in c76c65c, folded into the original fix. Production warning behavior is unchanged.

[kbattocchi]

The finite-OOB assertion is intentional for this fixed-seed test: it lets the reference use ordinary mean/std and the observation count instead of reproducing production's NaN-handling logic. Missing rows and unequal per-column coverage are exercised separately in test_cfdml_dr_stderr_per_column. If these seeds stop giving full OOB coverage, we'd rather update this fixture explicitly than silently change what the test checks.

[kbattocchi]

Agreed—making the convention explicit improves readability. I've added ddof=0 to both the per-column and fitted-estimator reference calculations. This preserves NumPy's existing default and does not change production behavior. The test-only follow-up is committed locally, pending publication approval.

[kbattocchi]

Published the explicit ddof=0 references in c76c65c, folded into the original fix.