uber / uber/causalml

sklearn metadata routing: let `treatment` reach a learner inside a Pipeline

Open
#986 7 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement
Dominant language
Python
Stars
6k
Forks
877
PR merge metrics
No merged PRs in 30d

Description

Part of #980 (v1.0 M1). Closes #854 when it lands.

Reordering fit to (X, y, treatment, …) is necessary but not sufficient for #854. Pipeline.fit(X, y) passes exactly two things; treatment is a third required array with nowhere to ride. After the flip a CausalML learner still can't be a Pipeline step without an adapter — the reordering just removes the first obstacle.

The mechanism sklearn provides is metadata routing: fit(X, y, **fit_params) with set_fit_request(treatment=True), so a caller writes pipe.fit(X, y, treatment=treatment) and the router delivers treatment to the step that asked for it.

Prerequisite

#985, the signature flip. Until then Pipeline's positional y binds to treatment, so an end-to-end test cannot be written. #985 is unblocked and targets v1.0 (Jun 2027).

Scope for the first pass

  • Meta-learners only — they are the classes users most want inside a Pipeline.
  • Integration tests and documentation before any custom MetadataRouter. BaseLearner subclasses BaseEstimator and the fit signatures declare treatment and p explicitly, so set_fit_request should be generated automatically; confirm that against sklearn.utils.metadata_routing before writing any routing code.
  • Cover an S-learner pipeline plus one propensity-consuming learner.
  • Cover both paths: routing enabled via sklearn.set_config(enable_metadata_routing=True), and the prefixed routing-disabled form.
  • Row-preserving transformers only. Resampling has to move X, y, treatment and p together and belongs in its own design.

Out of scope for the first pass: whether p should route into a learner's wrapped estimator, rather than being consumed by the learner itself.

Why it's separate

The argument order is a breaking change on a deadline (v1.0 API freeze); routing is additive and can land any time after. Coupling them would put a design-heavy feature on the critical path of a mechanical rename.

Acceptance

  • A documented, tested example of a CausalML learner as a Pipeline final step with treatment routed, no adapter class.
  • #854's original reproduction works.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with sklearn.utils.metadata_routing and Pipeline, then inspect BaseLearner and the meta-learner fit signatures after #985 lands. Add integration coverage for an S-learner and a propensity-consuming learner with enabled and prefixed routing, plus the documented Pipeline example. Done means treatment reaches the final learner without an adapter and #854's reproduction works.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, scikit-learn
Domain
documentation, machine-learning, testing-qa
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.