uber / uber/causalml

Adopt the scikit-learn `fit(X, y, treatment, ...)` argument order (v1.0 M1)

Open
#980 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

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

Description

Summary

CausalML learners take fit(X, treatment, y, ...), which puts y third. sklearn.pipeline.Pipeline calls the final estimator's fit(X, y) positionally, so today a CausalML learner can't be a Pipeline step without an adapter — the problem reported in #854.

The v1.0 roadmap RFC lists this under M1 — Foundation & stability as the one intentional breaking change: adopt fit(X, y, treatment, ...).

Reordering two required positional arguments can't be done in one step. y and treatment are both same-length arrays, so a straight swap would silently mis-train existing positional callers with no TypeError to catch it. So this is a two-step deprecation, and this epic tracks it end to end.

Step one has landed (#975): positional order is unchanged, but passing treatment/y positionally emits a FutureWarning pointing callers at keyword arguments — fit(X, y=y, treatment=treatment) — which are order-independent and therefore safe across the flip.

The rule

One rule, applied uniformly: X, then y, then treatment, and every other parameter keeps its relative position. It is implemented in causalml/inference/_arg_order.py as v1_order(), derived per method from that method's own signature, so the target order is machine-checkable rather than hand-maintained:

method today v1.0
meta-learner fit (X, treatment, y, p, …) (X, y, treatment, p, …)
predict (X, treatment, y, p, …) (X, y, treatment, p, …)
CausalTreeRegressor.fit (X, treatment, y, sample_weight, …) (X, y, treatment, sample_weight, …)
IVRegressor.fit (X, treatment, y, w) (X, y, treatment, w)
BaseDRIVLearner.fit (X, assignment, treatment, y, …) (X, y, treatment, assignment, …)

predict is included deliberately. Its current order doesn't block Pipeline, since predict(X) already works — but flipping only fit would leave v1.0 with the same two arguments in opposite positional orders on one class. Methods that don't take treatment/y at all, such as BaseRLearner.predict, are untouched.

Why the sequencing matters

A deprecation window is one-shot. Anything whose positional order changes at v1.0 has to warn in the same release as the shim, or it needs a second deprecation cycle of its own. That is why the "close the window" issues below are not optional cleanup — they gate the flip, and they must ship alongside #975 rather than with it.

The flip itself is inherently breaking and cannot be staged: once a signature is reordered, old positional calls break silently. It therefore lands as a single change on the v1.0 line, per-family checklist and all.

Child issues (dependency-ordered)

Close the deprecation window — ships with the warning release:

  • #981 — extend the shim to the remaining public methods that take treatment before y
  • #982 — migrate tests/ to keyword calls
  • #983 — migrate docs/ notebooks and guides to keyword calls
  • #984 — migration guide + docs/changelog.rst entry

The flip — v1.0:

  • #985 — reorder every signature to (X, y, treatment, …) and delete the shim

Actually deliver #854:

  • #986 — sklearn metadata routing so treatment can reach a learner inside a Pipeline

Reordering alone does not make a learner a drop-in Pipeline step: treatment still has to be threaded through sklearn's metadata routing. #854 stays open until that lands.

Open questions

All three are answered, and each answer is now pinned by shipped code rather than by this issue. The deprecation window is one-shot, so the FutureWarning users receive in 0.18.0 already names a target order for every affected method; #985 has to arrive at exactly those orders.

  1. UpliftTreeClassifier.fit takes (X, treatment, y, X_val, treatment_val, y_val, sample_weight, check_input). Reorder the *_val arguments too, or leave them and document the asymmetry? Answered: the suffixed pair is reordered in place, giving (X, y, treatment, X_val, y_val, treatment_val, sample_weight, check_input). Implemented as _order_suffixed_pairs() in causalml/inference/_arg_order.py (#981) and published in docs/migration.rst (#984).
  2. Sensitivity.get_prediction / .get_ate_ci / .get_potential_outcome_predictions take (X, p, treatment, y) — a third shape, with p second. In scope for the flip, or left alone as evaluation-side helpers? Answered: in scope. (X, p, treatment, y)(X, y, treatment, p). All three have been shimmed since #981 and appear in the migration guide's per-family table.
  3. Window length. How many minor releases between the warning and the flip? Answered: releases are quarterly — 0.18.0 Sep 2026 (ships the warning), 0.19.0 Dec 2026, 0.20.0 Mar 2027, v1.0 Jun 2027 (the flip). Three minor releases, roughly nine months. Published in the migration guide and changelog by #993, so #985 is no longer blocked.

Refs #854. Part of the v1.0 M1 milestone (#938).

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 child issue #985 and the ordering rules in causalml/inference/_arg_order.py. Review the affected method signatures and the per-family checklist, then inspect tests/ for positional-call coverage. Done means the v1.0 signatures follow the documented order and the deprecation shim is removed without leaving affected callers on the old order.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, scikit-learn
Domain
backend-api-design, machine-learning
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.