uber / uber/causalml

Flip every signature to `(X, y, treatment, ...)` and delete the deprecation shim

Open
#985 1 comment 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

Part of #980 (v1.0 M1). Refs #854.

The breaking change itself: reorder every signature to (X, y, treatment, …) and delete the deprecation shim.

This lands as one change. Once a signature is reordered, existing positional calls break silentlyy and treatment are both same-length arrays, so there is no TypeError to catch a swap. A partially flipped package is therefore worse than either end state, and no subset of this checklist may be released on its own.

Unblocked and scheduled. The window-closing issues are merged (#981 shim coverage, #982 tests, #983 docs, #984 migration guide), and the release schedule is fixed: the warning ships in 0.18.0 (Sep 2026), and this flip lands in v1.0 (Jun 2027) — three minor releases, roughly nine months.

The rule

X, then y, then treatment; every other parameter keeps its relative position. Already implemented as v1_order() in causalml/inference/_arg_order.py and derived per method from that method's own signature — so the target for each method is machine-readable, not hand-maintained. tests/test_fit_arg_order.py::test_v1_order pins it.

Two shapes that are not a plain swap of the leading pair, both settled and already published in docs/migration.rst:

  • A suffixed pair is reordered in place, so UpliftTreeClassifier.fit(X, y, treatment, X_val, y_val, treatment_val, sample_weight, check_input).
  • Sensitivity.get_* takes p second: (X, p, treatment, y)(X, y, treatment, p).

Scope

The authoritative list is every method carrying _arg_order_shimmed — 56 on a default install, plus 12 on the optional TF/torch/JAX backends. The window is one-shot, so each of those methods has already told users its v1.0 order, and each must arrive at exactly that order here. The checklist below is that set grouped by family:

  • Meta-learnersBaseSLearner, BaseTLearner, BaseXLearner, BaseRLearner, BaseDRLearner and subclasses (XGBTRegressor, XGBRRegressor, XGBRClassifier, XGBTClassifier, BaseSClassifier, BaseTClassifier, BaseXClassifier, BaseRClassifier, BaseDRClassifier, LRSRegressor), plus bootstrap / fit_bootstrap_ensemble
  • Causal trees and forestsCausalTreeRegressor (fit, fit_predict, estimate_ate, bootstrap, bootstrap_pool), CausalRandomForestRegressor
  • Uplift trees and forestsUpliftTreeClassifier / _KernelUpliftTreeClassifier (fit, plus fill and prune), UpliftRandomForestClassifier / _KernelUpliftRandomForestClassifier
  • IVIVRegressor.fit(X, y, treatment, w); BaseDRIVLearner.fit / fit_predict / estimate_ate / predict / bootstrap(X, y, treatment, assignment, …)
  • Sensitivity helpersSensitivity.get_prediction / get_ate_ci / get_potential_outcome_predictions(X, y, treatment, p)
  • Neural estimators — TF and JAX DragonNet, Torch and JAX CEVAE
  • StandalonePolicyLearner, TMLELearner
  • CausalML's own same-instance calls_uplift/upliftforest.py:342, _uplift/uplifttree.py:567 and meta/rlearner.py:913 call super().fit(X, treatment, y, …) positionally. The shim guards them today, so they warn nobody; once the order flips, treatment lands silently in y.
  • Update every Args: docstring block to match the new order
  • Delete causalml/inference/_arg_order.py, the __init_subclass__ hook in causalml/inference/serialization.py, and the six @shim_arg_order class decorators
  • Rewrite tests/test_fit_arg_order.py to assert the new order rather than the deprecation (the file's docstring says to update, not delete, these tests)
  • Move docs/migration.rst and the docs/changelog.rst deprecation entry to past tense — both currently describe the flip as forthcoming

Acceptance

  • Every method that carried _arg_order_shimmed before this change now has X, y, treatment leading inspect.signature(...).parameters, in the order v1_order() published for it.
  • No FutureWarning machinery remains; grep for _arg_order, shim_arg_order, _in_arg_order_call is empty.
  • Full suite green, and the equivalence test that guards against a silent y/treatment swap is retained in flipped form.

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 causalml/inference/_arg_order.py, tests/test_fit_arg_order.py, and the methods carrying _arg_order_shimmed; use v1_order() and inspect.signature to inventory the target order. Review the listed learner, tree, forest, IV, sensitivity, neural, and standalone entry points, plus the same-instance calls in the named files. Done means every listed method has the published order, the shim machinery is gone, updated tests and docs pass, and the full suite is green.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch, tensorflow
Domain
machine-learning
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.