Skip to content

Fix n_jobs not propagated to auto first-stage learners - #1045

Open
genrichez wants to merge 2 commits into
py-why:mainfrom
genrichez:fix/propagate-n-jobs-to-first-stage
Open

Fix n_jobs not propagated to auto first-stage learners#1045
genrichez wants to merge 2 commits into
py-why:mainfrom
genrichez:fix/propagate-n-jobs-to-first-stage

Conversation

@genrichez

Copy link
Copy Markdown

Summary

When using model_y='auto' or model_t='auto' in estimators like SparseLinearDML or CausalForestDML, the n_jobs parameter was not forwarded to the first-stage model selector. This meant that RandomForest and other parallelizable models always ran on a single core, regardless of the n_jobs setting passed to the estimator.
Fixes #1009

Changes

Added n_jobs parameter to get_selector( ) and _make_first_stage_selector(), and forwarded it from all estimator classes that support n_jobs:

File Change
econml/sklearn_extensions/model_selection.py get_selector() now accepts n_jobs and passes it to RF, GridSearchCV, and WeightedLassoCVWrapper
econml/dml/dml.py _make_first_stage_selector() accepts n_jobs; DML and NonParamDML forward it
econml/dml/causal_forest.py CausalForestDML forwards self.n_jobs
econml/dr/_drlearner.py DR learner's copy of _make_first_stage_selector() updated
econml/iv/dml/_dml.py IV/DML estimators forward n_jobs
econml/iv/dr/_dr.py IV/DR estimators forward n_jobs
econml/panel/dml/_dml.py DynamicDML forwards n_jobs

Classes without n_jobs (e.g. LinearDML) use getattr(self, 'n_jobs', None) so they default to None with no behavior change.
LogisticRegressionCV is intentionally skipped because sklearn 1.8 deprecated its n_jobs parameter.

How to verify

from econml.dml import SparseLinearDML

est = SparseLinearDML(n_jobs=4, random_state=0)
model_y = est._gen_model_y()

# Before: RF inside the selector has n_jobs=None
# After: RF inside the selector has n_jobs=4
selector = model_y._model
for m in selector.models:
    if hasattr(m, 'searcher') and hasattr(m.searcher, 'estimator'):
        print(f"{type(m.searcher.estimator).__name__}.n_jobs = {m.searcher.estimator.n_jobs}")

Tests

All existing tests pass (excluding unrelated ray tests that require the ray package)

genrichez and others added 2 commits July 5, 2026 13:32
When using model_y='auto' or model_t='auto', the n_jobs parameter was
not forwarded to the first-stage model selector. This meant that
RandomForest and other parallelizable models in the auto selector
always ran on a single core regardless of the n_jobs setting.

Add n_jobs parameter to get_selector() and _make_first_stage_selector(),
and forward it from all DML, DR, and IV estimator classes that have
n_jobs. Classes without n_jobs (e.g. LinearDML) default to None.

Skip LogisticRegressionCV since sklearn 1.8 deprecated its n_jobs param.

Fixes microsoft#1009

Signed-off-by: genrichez <2.2434764e+07+genrichez@users.noreply.github.com>
@genrichez

Copy link
Copy Markdown
Author

Hi @kbattocchi following up on this one since it's been quiet a few weeks. You folded my other PR #1046 into #1047 recently, so figured I'd flag this one too in case it's useful, fixes n_jobs not propagating to the auto first-stage learners. Happy to adjust if you'd like a different approach.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

n_jobs not passed to 'auto' learners in SparseLinearDML()

1 participant