Skip to content

Deprecate positional treatment/y in meta-learner fit() for sklearn Pipeline compat (#854) - #975

Open
jeongyoonlee wants to merge 1 commit into
masterfrom
feature/854-fit-arg-order
Open

Deprecate positional treatment/y in meta-learner fit() for sklearn Pipeline compat (#854)#975
jeongyoonlee wants to merge 1 commit into
masterfrom
feature/854-fit-arg-order

Conversation

@jeongyoonlee

Copy link
Copy Markdown
Collaborator

What & why

Kicks off #854 — making the meta-learners compatible with sklearn.pipeline.Pipeline, the one intentional breaking change on the v1.0 roadmap (M1).

Pipeline calls the final estimator's fit(X, y) positionally, but CausalML meta-learners take fit(X, treatment, y, ...)y is third. In v1.0 the positional order becomes fit(X, y, treatment, ...).

Reordering two required positional arguments can't be done in one step without risking a silent y/treatment swap for existing positional callers. So this is a two-step deprecation, and this PR is step one (non-breaking):

  • Positional order is unchanged. Every existing fit(X, treatment, y) call keeps working exactly as before.
  • Passing treatment/y positionally now emits a FutureWarning pointing callers to keyword arguments: fit(X, y=y, treatment=treatment). Keyword calls are order-independent, so code that migrates now keeps working unchanged across the v1.0 flip.
  • Step two (a later PR, targeted at v1.0): flip the positional signature to (X, y, treatment, ...) and drop the shim.

How

  • A small re-entrancy-guarded decorator in BaseLearner, wired in via __init_subclass__, wraps each subclass's own fit / fit_predict / estimate_ate. This covers the whole S/T/X/R/DR family (and subclasses like XGBTRegressor) with no per-method edits, and functools.wraps preserves the real signatures for introspection.
  • The guard ensures internal delegation — fit_predict → fit, subclass fit → super().fit(X, treatment, y, *args, **kwargs), estimate_ate → fit_predict → fit — warns at most once per top-level call.
  • No internal call sites needed changing: CausalML's own code (metrics/sensitivity.py, dataset/synthetic.py, …) already calls these methods with keywords.

Scope

  • In: the BaseLearner meta-learner family (S/T/X/R/DR + subclasses).
  • Out: predict already takes X first with treatment/y optional, so Pipeline's predict(X) already works — no change needed. TMLELearner (standalone, not BaseLearner, not Pipeline-usable) and the IV/NN learners are untouched here.
  • Note: reordering alone doesn't make a meta-learner a drop-in Pipeline step — treatment still needs to be threaded via sklearn metadata routing. That's a separate follow-up; this PR is the argument-order groundwork.

Tests

tests/test_fit_arg_order.py (19 cases): positional-use warning across S/T/X/R/DR, keyword silence, single-warning delegation (fit_predict / estimate_ate / subclass super().fit), signature preservation, and positional == keyword equivalence (guards against a silent y/treatment swap). Full tests/test_meta_learners.py (74) stays green.

Refs #854.

Kick off the scikit-learn Pipeline-compatibility change for meta-learners:
the positional fit() argument order will move from (X, treatment, y, ...) to
(X, y, treatment, ...) in v1.0. This is the first, non-breaking step of a
two-step deprecation.

- Add a re-entrancy-guarded shim in BaseLearner that emits a FutureWarning
  when treatment/y are passed positionally to fit/fit_predict/estimate_ate,
  steering callers to order-independent keyword arguments. It is wired in via
  __init_subclass__ so every S/T/X/R/DR learner (and its subclasses) is
  covered with no per-method edits, and the real signatures are preserved via
  functools.wraps. The guard ensures internal delegation (fit_predict -> fit,
  subclass fit -> super().fit, estimate_ate -> fit_predict -> fit) warns at
  most once per top-level call.
- The positional order is UNCHANGED for now, so no existing call silently
  breaks; keyword calls -- fit(X, y=y, treatment=treatment) -- are silent and
  stay correct across the v1.0 flip.
- Add tests/test_fit_arg_order.py covering the warning on positional use,
  keyword silence, single-warning delegation, subclass super().fit,
  signature preservation, and positional==keyword equivalence (guarding
  against a silent y/treatment swap).
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.

2 participants