From b9a675117cf9deda3adbdd8cf31d8774d979e105 Mon Sep 17 00:00:00 2001
From: Jeong-Yoon Lee <jeongyoon.lee1@gmail.com>
Date: Sun, 26 Jul 2026 11:21:06 -0700
Subject: [PATCH 1/3] Deprecate positional treatment/y in meta-learner fit()
 (#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).
---
 causalml/inference/meta/base.py |  67 ++++++++++++++++++
 tests/test_fit_arg_order.py     | 120 ++++++++++++++++++++++++++++++++
 2 files changed, 187 insertions(+)
 create mode 100644 tests/test_fit_arg_order.py

diff --git a/causalml/inference/meta/base.py b/causalml/inference/meta/base.py
index c8792166..9ce5cad9 100644
--- a/causalml/inference/meta/base.py
+++ b/causalml/inference/meta/base.py
@@ -1,5 +1,8 @@
 from abc import ABCMeta, abstractmethod
+import functools
 import logging
+import warnings
+
 import numpy as np
 import pandas as pd
 from joblib import Parallel, delayed
@@ -48,6 +51,56 @@ def _fit_bootstrap_clone(learner_template, X, treatment, y, p, seed, bootstrap_s
     return learner_b
 
 
+# --- #854: scikit-learn `fit(X, y, ...)` argument-order migration ------------
+# CausalML meta-learners historically take ``fit(X, treatment, y, ...)``, which
+# puts ``y`` third and breaks ``sklearn.pipeline.Pipeline`` (it calls the final
+# estimator's ``fit(X, y)`` positionally). In v1.0 the positional order becomes
+# ``fit(X, y, treatment, ...)``. This is a two-step deprecation: for now the
+# positional order is UNCHANGED (no silent breakage), but passing ``treatment``
+# or ``y`` positionally emits a ``FutureWarning`` steering callers to keyword
+# arguments -- ``fit(X, y=y, treatment=treatment)`` -- which are order-independent
+# and therefore safe across the v1.0 flip. Remove this shim when the signatures
+# are reordered for v1.0.
+_FIT_ARG_ORDER_MSG = (
+    "Passing `treatment` and/or `y` to {method}() by position is deprecated and "
+    "will change in causalml v1.0: the positional argument order will switch from "
+    "(X, treatment, y, ...) to (X, y, treatment, ...) for scikit-learn Pipeline "
+    "compatibility (see https://github.com/uber/causalml/issues/854). To be safe "
+    "across the change, pass them as keyword arguments, e.g. "
+    "{method}(X, y=y, treatment=treatment)."
+)
+
+
+def _deprecate_positional_treatment_y(method):
+    """Warn when ``treatment``/``y`` are passed positionally to a fit-family method.
+
+    Only the outermost fit-family call on a given instance warns, so internal
+    delegations (``fit_predict`` -> ``fit``, a subclass ``fit`` -> ``super().fit``,
+    ``estimate_ate`` -> ``fit_predict`` -> ``fit``) never double-count and callers
+    see at most one warning per top-level call. The wrapped method keeps its real
+    signature via :func:`functools.wraps`, so introspection is unaffected.
+    """
+
+    @functools.wraps(method)
+    def wrapper(self, X, *args, **kwargs):
+        if getattr(self, "_in_fit_family_call", False):
+            return method(self, X, *args, **kwargs)
+        if args:
+            warnings.warn(
+                _FIT_ARG_ORDER_MSG.format(method=method.__name__),
+                FutureWarning,
+                stacklevel=2,
+            )
+        self._in_fit_family_call = True
+        try:
+            return method(self, X, *args, **kwargs)
+        finally:
+            self._in_fit_family_call = False
+
+    wrapper._treatment_y_shimmed = True
+    return wrapper
+
+
 class BaseLearner(SerializableLearner, BaseEstimator, metaclass=ABCMeta):
     """Base class for all causalml meta-learners.
 
@@ -68,6 +121,19 @@ class BaseLearner(SerializableLearner, BaseEstimator, metaclass=ABCMeta):
     * ``__repr__`` is inherited from ``BaseEstimator`` and reflects constructor params.
     """
 
+    def __init_subclass__(cls, **kwargs):
+        """Wrap each subclass's own fit-family methods with the #854 arg-order shim.
+
+        Only methods defined directly on ``cls`` are wrapped (inherited ones are
+        already wrapped on the parent), and already-wrapped methods are skipped,
+        so no method is double-wrapped.
+        """
+        super().__init_subclass__(**kwargs)
+        for name in ("fit", "fit_predict", "estimate_ate"):
+            method = cls.__dict__.get(name)
+            if callable(method) and not getattr(method, "_treatment_y_shimmed", False):
+                setattr(cls, name, _deprecate_positional_treatment_y(method))
+
     def _is_fitted(self):
         """Meta-learners are fitted once t_groups is set during fit()."""
         return hasattr(self, "t_groups")
@@ -84,6 +150,7 @@ def predict(
     ):
         pass
 
+    @_deprecate_positional_treatment_y
     def fit_predict(
         self,
         X,
diff --git a/tests/test_fit_arg_order.py b/tests/test_fit_arg_order.py
new file mode 100644
index 00000000..ffbbd7be
--- /dev/null
+++ b/tests/test_fit_arg_order.py
@@ -0,0 +1,120 @@
+"""Tests for the #854 fit-family argument-order deprecation shim.
+
+CausalML meta-learners historically take ``fit(X, treatment, y, ...)``, which puts
+``y`` third and breaks ``sklearn.pipeline.Pipeline`` (it calls the final
+estimator's ``fit(X, y)`` positionally). For scikit-learn compatibility the
+positional order becomes ``fit(X, y, treatment, ...)`` in v1.0.
+
+This is a two-step deprecation: the positional order is UNCHANGED for now (so no
+existing call silently breaks), but passing ``treatment``/``y`` positionally emits
+a ``FutureWarning`` steering callers to keyword arguments, which are
+order-independent and therefore safe across the v1.0 flip. These tests pin that
+behavior; they should be updated (not deleted) when the signatures are reordered.
+"""
+
+import inspect
+import warnings
+
+import numpy as np
+import pytest
+from sklearn.linear_model import LinearRegression
+
+from causalml.inference.meta import (
+    BaseSRegressor,
+    BaseTRegressor,
+    BaseXRegressor,
+    BaseRRegressor,
+    BaseDRRegressor,
+    XGBTRegressor,
+)
+
+# One representative regressor per meta-learner family; all default to
+# ``control_name=0`` which matches the 0/1 treatment from ``synthetic_data``.
+REGRESSORS = [
+    BaseSRegressor,
+    BaseTRegressor,
+    BaseXRegressor,
+    BaseRRegressor,
+    BaseDRRegressor,
+]
+
+
+def _order_warnings(record):
+    """Filter a warning record down to the #854 arg-order FutureWarnings."""
+    return [
+        w
+        for w in record
+        if issubclass(w.category, FutureWarning) and "argument order" in str(w.message)
+    ]
+
+
+@pytest.mark.parametrize("learner_cls", REGRESSORS)
+def test_positional_fit_warns(generate_regression_data, learner_cls):
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = learner_cls(learner=LinearRegression())
+    with pytest.warns(FutureWarning, match="argument order"):
+        learner.fit(X, treatment, y)
+
+
+@pytest.mark.parametrize("learner_cls", REGRESSORS)
+def test_keyword_fit_does_not_warn(generate_regression_data, learner_cls):
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = learner_cls(learner=LinearRegression())
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        learner.fit(X=X, y=y, treatment=treatment)
+    assert _order_warnings(record) == []
+
+
+def test_positional_and_keyword_are_equivalent(generate_regression_data):
+    """The shim must never silently swap y/treatment: both call styles must fit
+    identically. Guards against a values-based reorder that could corrupt data."""
+    y, X, treatment, _, _, _ = generate_regression_data()
+
+    positional = BaseTRegressor(learner=LinearRegression())
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", FutureWarning)
+        positional.fit(X, treatment, y)
+
+    keyword = BaseTRegressor(learner=LinearRegression())
+    keyword.fit(X=X, y=y, treatment=treatment)
+
+    np.testing.assert_allclose(positional.predict(X), keyword.predict(X))
+
+
+def test_fit_predict_warns_once(generate_regression_data):
+    """fit_predict -> fit delegation must not double-count the warning."""
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = BaseTRegressor(learner=LinearRegression())
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        learner.fit_predict(X, treatment, y)
+    assert len(_order_warnings(record)) == 1
+
+
+def test_estimate_ate_warns_once(generate_regression_data):
+    """estimate_ate -> (fit_predict ->) fit delegation must warn exactly once."""
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = BaseSRegressor(learner=LinearRegression())
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        learner.estimate_ate(X, treatment, y)
+    assert len(_order_warnings(record)) == 1
+
+
+def test_subclass_super_fit_warns_once(generate_regression_data):
+    """A subclass fit that delegates to super().fit (XGBTRegressor) warns once."""
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = XGBTRegressor()
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        learner.fit(X, treatment, y)
+    assert len(_order_warnings(record)) == 1
+
+
+@pytest.mark.parametrize("learner_cls", REGRESSORS)
+def test_fit_signature_is_preserved(learner_cls):
+    """functools.wraps keeps the real (X, treatment, y, ...) signature visible."""
+    learner = learner_cls(learner=LinearRegression())
+    params = list(inspect.signature(learner.fit).parameters)
+    assert params[:3] == ["X", "treatment", "y"]

From 88cde2f7d410e47d9f805702e39595f8fe40d75b Mon Sep 17 00:00:00 2001
From: Jeong-Yoon Lee <jeongyoon.lee1@gmail.com>
Date: Fri, 31 Jul 2026 11:08:15 +0900
Subject: [PATCH 2/3] Extend the #854 arg-order deprecation to every learner
 family

Step one of #854 warned only on the BaseLearner meta-learners. A
deprecation window is one-shot: anything whose positional order changes
at v1.0 has to warn in this release or it needs a second cycle. Widen
the shim so v1.0 can flip the whole package at once.

- Move the shim to causalml/inference/_arg_order.py and derive the
  target order from each method's own signature: X, then y, then
  treatment, with every other parameter keeping its relative position.
  Each warning now names that method's real order instead of assuming
  (X, treatment, y, ...). BaseDRIVLearner.fit therefore resolves to
  (X, y, treatment, assignment, ...) and IVRegressor.fit to
  (X, y, treatment, w).

- Host __init_subclass__ on SerializableLearner rather than BaseLearner.
  That one move covers the causal/uplift trees and forests and the IV
  learners as well as the meta-learners. The trees matter most: they
  already inherit sklearn's ForestRegressor/BaseDecisionTree and
  override fit(X, y, ...) with fit(X, treatment, y, ...), so leaving
  them behind would ship v1.0 with the most sklearn-shaped classes
  keeping the least sklearn-shaped signature.

- Decorate the six learners outside that hierarchy: the TF/JAX
  DragonNet, the Torch/JAX CEVAE, PolicyLearner and TMLELearner.

- Shim predict too. Its (X, treatment, y, ...) is unaffected by
  Pipeline, but without this v1.0 would leave fit and predict holding
  the same two arguments in opposite orders on one class. Methods
  without treatment/y, such as BaseRLearner.predict, are returned
  unwrapped.

Tests go 19 -> 43, including a completeness guard that walks the
package and fails if any method taking treatment and y positionally
lacks the shim, so a learner added later cannot miss the window.

Refs #854.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 causalml/inference/_arg_order.py      | 124 ++++++++++++++++
 causalml/inference/jax/cevae/cevae.py |   2 +
 causalml/inference/jax/dragonnet.py   |   2 +
 causalml/inference/meta/base.py       |  66 ---------
 causalml/inference/meta/tmle.py       |   2 +
 causalml/inference/serialization.py   |  13 ++
 causalml/inference/tf/dragonnet.py    |   2 +
 causalml/inference/torch/cevae.py     |   2 +
 causalml/optimize/policylearner.py    |   2 +
 tests/test_fit_arg_order.py           | 198 +++++++++++++++++++++++++-
 10 files changed, 344 insertions(+), 69 deletions(-)
 create mode 100644 causalml/inference/_arg_order.py

diff --git a/causalml/inference/_arg_order.py b/causalml/inference/_arg_order.py
new file mode 100644
index 00000000..f5ee5694
--- /dev/null
+++ b/causalml/inference/_arg_order.py
@@ -0,0 +1,124 @@
+"""#854: scikit-learn ``fit(X, y, ...)`` argument-order migration.
+
+CausalML learners historically take ``fit(X, treatment, y, ...)``, which puts
+``y`` third and breaks ``sklearn.pipeline.Pipeline`` (it calls the final
+estimator's ``fit(X, y)`` positionally). In v1.0 the positional order becomes
+``(X, y, treatment, ...)``; every other parameter keeps its relative position,
+so ``BaseDRIVLearner.fit(X, assignment, treatment, y, ...)`` becomes
+``fit(X, y, treatment, assignment, ...)``.
+
+This is a two-step deprecation. Step one (this module): the positional order is
+UNCHANGED, so no existing call silently breaks, but passing ``treatment`` or
+``y`` positionally emits a ``FutureWarning`` steering callers to keyword
+arguments, which are order-independent and therefore safe across the flip.
+Step two (v1.0): reorder the signatures and delete this module.
+"""
+
+import functools
+import inspect
+import warnings
+
+#: Parameters whose positional slot changes in v1.0, in their v1.0 order.
+REORDERED_PARAMS = ("y", "treatment")
+
+#: Methods wrapped automatically for ``SerializableLearner`` subclasses.
+SHIMMED_METHODS = ("fit", "fit_predict", "estimate_ate", "predict")
+
+_MSG = (
+    "Passing `treatment` and/or `y` to {name}() by position is deprecated and "
+    "will change in causalml v1.0: the positional argument order becomes "
+    "({new_order}) for scikit-learn Pipeline compatibility "
+    "(see https://github.com/uber/causalml/issues/854). To be safe across the "
+    "change, pass them as keyword arguments, e.g. {example}."
+)
+
+
+def _positional_params(method):
+    """Return the names of ``method``'s positional parameters, minus ``self``."""
+    kinds = (
+        inspect.Parameter.POSITIONAL_ONLY,
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+    )
+    names = [
+        p.name for p in inspect.signature(method).parameters.values() if p.kind in kinds
+    ]
+    return names[1:] if names[:1] == ["self"] else names
+
+
+def v1_order(params):
+    """Reorder ``params`` to the v1.0 convention: X, then y, treatment, then rest.
+
+    Parameters other than ``y``/``treatment`` keep their relative order, so a
+    signature's extra arguments (``assignment``, ``p``, ``sample_weight``, ...)
+    simply shift right.
+    """
+    head, tail = params[:1], params[1:]
+    moved = [name for name in REORDERED_PARAMS if name in tail]
+    return head + moved + [name for name in tail if name not in REORDERED_PARAMS]
+
+
+def deprecate_positional_treatment_y(method):
+    """Warn when ``treatment``/``y`` are passed positionally to ``method``.
+
+    ``method`` is returned unchanged when neither parameter is positional, so
+    this is safe to apply blanket-style from ``__init_subclass__``.
+
+    Only the outermost shimmed call on a given instance warns, so internal
+    delegation (``fit_predict`` -> ``fit``/``predict``, a subclass ``fit`` ->
+    ``super().fit``, ``estimate_ate`` -> ``fit_predict`` -> ``fit``) never
+    double-counts and callers see at most one warning per top-level call.
+    :func:`functools.wraps` preserves the real signature, so introspection and
+    scikit-learn metadata routing are unaffected.
+    """
+    params = _positional_params(method)
+    if not any(name in REORDERED_PARAMS for name in params):
+        return method
+
+    message = _MSG.format(
+        name=method.__name__,
+        new_order=", ".join(v1_order(params)),
+        example="{}({}, {})".format(
+            method.__name__,
+            params[0],
+            ", ".join(f"{p}={p}" for p in REORDERED_PARAMS if p in params),
+        ),
+    )
+
+    @functools.wraps(method)
+    def wrapper(self, *args, **kwargs):
+        if getattr(self, "_in_arg_order_call", False):
+            return method(self, *args, **kwargs)
+        if any(
+            params[i] in REORDERED_PARAMS for i in range(min(len(args), len(params)))
+        ):
+            warnings.warn(message, FutureWarning, stacklevel=2)
+        self._in_arg_order_call = True
+        try:
+            return method(self, *args, **kwargs)
+        finally:
+            self._in_arg_order_call = False
+
+    wrapper._arg_order_shimmed = True
+    return wrapper
+
+
+def shim_arg_order(cls):
+    """Wrap ``cls``'s own :data:`SHIMMED_METHODS` with the arg-order shim.
+
+    Only methods defined directly on ``cls`` are wrapped (inherited ones are
+    already wrapped on the parent), abstract methods are left alone, and
+    already-wrapped methods are skipped, so nothing is double-wrapped.
+
+    ``SerializableLearner`` applies this automatically via ``__init_subclass__``.
+    Learners outside that hierarchy — the TF/Torch/JAX estimators,
+    ``PolicyLearner``, ``TMLELearner`` — use it as a class decorator.
+    """
+    for name in SHIMMED_METHODS:
+        method = cls.__dict__.get(name)
+        if (
+            callable(method)
+            and not getattr(method, "_arg_order_shimmed", False)
+            and not getattr(method, "__isabstractmethod__", False)
+        ):
+            setattr(cls, name, deprecate_positional_treatment_y(method))
+    return cls
diff --git a/causalml/inference/jax/cevae/cevae.py b/causalml/inference/jax/cevae/cevae.py
index 43ff8d1c..cb75efb9 100644
--- a/causalml/inference/jax/cevae/cevae.py
+++ b/causalml/inference/jax/cevae/cevae.py
@@ -24,6 +24,7 @@
 import optax
 
 from causalml.inference.meta.utils import convert_pd_to_np
+from causalml.inference._arg_order import shim_arg_order
 from causalml.inference.jax.cevae.losses import cevae_loss
 from causalml.inference.jax.cevae.modeling import Guide, Model, PreWhitener
 
@@ -80,6 +81,7 @@ def _loss(net):
     return train_step
 
 
+@shim_arg_order
 class CEVAE:
     """JAX/flax.nnx CEVAE for treatment-effect estimation.
 
diff --git a/causalml/inference/jax/dragonnet.py b/causalml/inference/jax/dragonnet.py
index 0b53335e..ab4f7607 100644
--- a/causalml/inference/jax/dragonnet.py
+++ b/causalml/inference/jax/dragonnet.py
@@ -25,6 +25,7 @@
     make_tarreg_loss,
 )
 from causalml.inference.meta.utils import convert_pd_to_np
+from causalml.inference._arg_order import shim_arg_order
 
 
 class EpsilonLayer(nnx.Module):
@@ -278,6 +279,7 @@ def _run_training_loop(
             break
 
 
+@shim_arg_order
 class DragonNet:
     """JAX/flax.nnx DragonNet for treatment effect estimation.
 
diff --git a/causalml/inference/meta/base.py b/causalml/inference/meta/base.py
index 9ce5cad9..e50d40cb 100644
--- a/causalml/inference/meta/base.py
+++ b/causalml/inference/meta/base.py
@@ -1,7 +1,5 @@
 from abc import ABCMeta, abstractmethod
-import functools
 import logging
-import warnings
 
 import numpy as np
 import pandas as pd
@@ -51,56 +49,6 @@ def _fit_bootstrap_clone(learner_template, X, treatment, y, p, seed, bootstrap_s
     return learner_b
 
 
-# --- #854: scikit-learn `fit(X, y, ...)` argument-order migration ------------
-# CausalML meta-learners historically take ``fit(X, treatment, y, ...)``, which
-# puts ``y`` third and breaks ``sklearn.pipeline.Pipeline`` (it calls the final
-# estimator's ``fit(X, y)`` positionally). In v1.0 the positional order becomes
-# ``fit(X, y, treatment, ...)``. This is a two-step deprecation: for now the
-# positional order is UNCHANGED (no silent breakage), but passing ``treatment``
-# or ``y`` positionally emits a ``FutureWarning`` steering callers to keyword
-# arguments -- ``fit(X, y=y, treatment=treatment)`` -- which are order-independent
-# and therefore safe across the v1.0 flip. Remove this shim when the signatures
-# are reordered for v1.0.
-_FIT_ARG_ORDER_MSG = (
-    "Passing `treatment` and/or `y` to {method}() by position is deprecated and "
-    "will change in causalml v1.0: the positional argument order will switch from "
-    "(X, treatment, y, ...) to (X, y, treatment, ...) for scikit-learn Pipeline "
-    "compatibility (see https://github.com/uber/causalml/issues/854). To be safe "
-    "across the change, pass them as keyword arguments, e.g. "
-    "{method}(X, y=y, treatment=treatment)."
-)
-
-
-def _deprecate_positional_treatment_y(method):
-    """Warn when ``treatment``/``y`` are passed positionally to a fit-family method.
-
-    Only the outermost fit-family call on a given instance warns, so internal
-    delegations (``fit_predict`` -> ``fit``, a subclass ``fit`` -> ``super().fit``,
-    ``estimate_ate`` -> ``fit_predict`` -> ``fit``) never double-count and callers
-    see at most one warning per top-level call. The wrapped method keeps its real
-    signature via :func:`functools.wraps`, so introspection is unaffected.
-    """
-
-    @functools.wraps(method)
-    def wrapper(self, X, *args, **kwargs):
-        if getattr(self, "_in_fit_family_call", False):
-            return method(self, X, *args, **kwargs)
-        if args:
-            warnings.warn(
-                _FIT_ARG_ORDER_MSG.format(method=method.__name__),
-                FutureWarning,
-                stacklevel=2,
-            )
-        self._in_fit_family_call = True
-        try:
-            return method(self, X, *args, **kwargs)
-        finally:
-            self._in_fit_family_call = False
-
-    wrapper._treatment_y_shimmed = True
-    return wrapper
-
-
 class BaseLearner(SerializableLearner, BaseEstimator, metaclass=ABCMeta):
     """Base class for all causalml meta-learners.
 
@@ -121,19 +69,6 @@ class BaseLearner(SerializableLearner, BaseEstimator, metaclass=ABCMeta):
     * ``__repr__`` is inherited from ``BaseEstimator`` and reflects constructor params.
     """
 
-    def __init_subclass__(cls, **kwargs):
-        """Wrap each subclass's own fit-family methods with the #854 arg-order shim.
-
-        Only methods defined directly on ``cls`` are wrapped (inherited ones are
-        already wrapped on the parent), and already-wrapped methods are skipped,
-        so no method is double-wrapped.
-        """
-        super().__init_subclass__(**kwargs)
-        for name in ("fit", "fit_predict", "estimate_ate"):
-            method = cls.__dict__.get(name)
-            if callable(method) and not getattr(method, "_treatment_y_shimmed", False):
-                setattr(cls, name, _deprecate_positional_treatment_y(method))
-
     def _is_fitted(self):
         """Meta-learners are fitted once t_groups is set during fit()."""
         return hasattr(self, "t_groups")
@@ -150,7 +85,6 @@ def predict(
     ):
         pass
 
-    @_deprecate_positional_treatment_y
     def fit_predict(
         self,
         X,
diff --git a/causalml/inference/meta/tmle.py b/causalml/inference/meta/tmle.py
index 372d0f3d..d1363c45 100644
--- a/causalml/inference/meta/tmle.py
+++ b/causalml/inference/meta/tmle.py
@@ -11,6 +11,7 @@
     check_p_conditions,
     convert_pd_to_np,
 )
+from causalml.inference._arg_order import shim_arg_order
 
 logger = logging.getLogger("causalml")
 
@@ -91,6 +92,7 @@ def simple_tmle(y, w, q0w, q1w, p, alpha=0.0001):
     return np.mean(q1star - q0star), np.sqrt(np.var(ic) / np.size(y))
 
 
+@shim_arg_order
 class TMLELearner:
     """Targeted maximum likelihood estimation.
 
diff --git a/causalml/inference/serialization.py b/causalml/inference/serialization.py
index cdcf8bc3..1bfc8557 100644
--- a/causalml/inference/serialization.py
+++ b/causalml/inference/serialization.py
@@ -14,6 +14,8 @@
 
 import joblib
 
+from causalml.inference._arg_order import shim_arg_order
+
 logger = logging.getLogger("causalml")
 
 
@@ -44,6 +46,17 @@ class SerializableLearner:
     check does not apply (e.g. for non-sklearn learners).
     """
 
+    def __init_subclass__(cls, **kwargs):
+        """Apply the #854 argument-order shim to each subclass.
+
+        Hosting this on the shared mixin rather than on a single learner base
+        covers the meta-learners, the causal/uplift trees and forests, and the
+        IV learners in one place. See :func:`~causalml.inference._arg_order.
+        shim_arg_order` for what gets wrapped.
+        """
+        super().__init_subclass__(**kwargs)
+        shim_arg_order(cls)
+
     def _is_fitted(self):
         """Check whether this learner has been fitted.
 
diff --git a/causalml/inference/tf/dragonnet.py b/causalml/inference/tf/dragonnet.py
index f8cd3ed3..72ae4c75 100644
--- a/causalml/inference/tf/dragonnet.py
+++ b/causalml/inference/tf/dragonnet.py
@@ -33,8 +33,10 @@
     make_tarreg_loss,
 )
 from causalml.inference.meta.utils import convert_pd_to_np
+from causalml.inference._arg_order import shim_arg_order
 
 
+@shim_arg_order
 class DragonNet:
     def __init__(
         self,
diff --git a/causalml/inference/torch/cevae.py b/causalml/inference/torch/cevae.py
index 3e7b0f79..68d6e488 100644
--- a/causalml/inference/torch/cevae.py
+++ b/causalml/inference/torch/cevae.py
@@ -27,6 +27,7 @@
 from pyro.contrib.cevae import CEVAE as CEVAEModel
 
 from causalml.inference.meta.utils import convert_pd_to_np
+from causalml.inference._arg_order import shim_arg_order
 
 pyro_logger = logging.getLogger("pyro")
 pyro_logger.setLevel(logging.DEBUG)
@@ -34,6 +35,7 @@
     pyro_logger.handlers[0].setLevel(logging.DEBUG)
 
 
+@shim_arg_order
 class CEVAE:
     def __init__(
         self,
diff --git a/causalml/optimize/policylearner.py b/causalml/optimize/policylearner.py
index 1dea008a..349cc669 100644
--- a/causalml/optimize/policylearner.py
+++ b/causalml/optimize/policylearner.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 from causalml.propensity import compute_propensity_score
+from causalml.inference._arg_order import shim_arg_order
 from sklearn.ensemble import GradientBoostingRegressor, GradientBoostingClassifier
 from sklearn.model_selection import KFold
 from sklearn.tree import DecisionTreeClassifier
@@ -9,6 +10,7 @@
 logger = logging.getLogger("causalml")
 
 
+@shim_arg_order
 class PolicyLearner:
     """
     A Learner that learns a treatment assignment policy with observational data using doubly robust estimator of causal
diff --git a/tests/test_fit_arg_order.py b/tests/test_fit_arg_order.py
index ffbbd7be..00d918d2 100644
--- a/tests/test_fit_arg_order.py
+++ b/tests/test_fit_arg_order.py
@@ -1,9 +1,12 @@
-"""Tests for the #854 fit-family argument-order deprecation shim.
+"""Tests for the #854 argument-order deprecation shim.
 
-CausalML meta-learners historically take ``fit(X, treatment, y, ...)``, which puts
+CausalML learners historically take ``fit(X, treatment, y, ...)``, which puts
 ``y`` third and breaks ``sklearn.pipeline.Pipeline`` (it calls the final
 estimator's ``fit(X, y)`` positionally). For scikit-learn compatibility the
-positional order becomes ``fit(X, y, treatment, ...)`` in v1.0.
+positional order becomes ``(X, y, treatment, ...)`` in v1.0 — across the whole
+package (meta-learners, causal/uplift trees and forests, IV, and the TF/Torch/
+JAX estimators), and for ``predict`` as well as the fit family, so the same two
+arguments never sit in opposite orders on one class.
 
 This is a two-step deprecation: the positional order is UNCHANGED for now (so no
 existing call silently breaks), but passing ``treatment``/``y`` positionally emits
@@ -12,13 +15,21 @@
 behavior; they should be updated (not deleted) when the signatures are reordered.
 """
 
+import importlib
 import inspect
+import pkgutil
 import warnings
 
 import numpy as np
 import pytest
 from sklearn.linear_model import LinearRegression
 
+from causalml.inference._arg_order import (
+    SHIMMED_METHODS,
+    _positional_params,
+    v1_order,
+)
+from causalml.inference.iv import BaseDRIVRegressor, IVRegressor
 from causalml.inference.meta import (
     BaseSRegressor,
     BaseTRegressor,
@@ -26,7 +37,14 @@
     BaseRRegressor,
     BaseDRRegressor,
     XGBTRegressor,
+    BaseSClassifier,
+    BaseTClassifier,
+    BaseXClassifier,
+    BaseRClassifier,
 )
+from causalml.inference.tree import CausalTreeRegressor
+
+from .const import RANDOM_SEED
 
 # One representative regressor per meta-learner family; all default to
 # ``control_name=0`` which matches the 0/1 treatment from ``synthetic_data``.
@@ -38,6 +56,16 @@
     BaseDRRegressor,
 ]
 
+# The R-learner's ``predict(X, p, return_components)`` already omits treatment/y,
+# so it is unaffected by the flip. Derive the list instead of hard-coding it, so
+# it stays correct if a predict signature changes.
+PREDICT_REORDERED = [
+    cls
+    for cls in REGRESSORS
+    if "treatment"
+    in _positional_params(getattr(cls.predict, "__wrapped__", cls.predict))
+]
+
 
 def _order_warnings(record):
     """Filter a warning record down to the #854 arg-order FutureWarnings."""
@@ -118,3 +146,167 @@ def test_fit_signature_is_preserved(learner_cls):
     learner = learner_cls(learner=LinearRegression())
     params = list(inspect.signature(learner.fit).parameters)
     assert params[:3] == ["X", "treatment", "y"]
+
+
+# --- v1.0 target ordering ---------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "current, expected",
+    [
+        # Meta-learners: treatment and y simply swap.
+        (["X", "treatment", "y", "p"], ["X", "y", "treatment", "p"]),
+        # predict: same rule, so fit and predict agree after the flip.
+        (
+            ["X", "treatment", "y", "p", "return_components", "verbose"],
+            ["X", "y", "treatment", "p", "return_components", "verbose"],
+        ),
+        # IV: the instrument keeps its relative position and shifts right.
+        (["X", "treatment", "y", "w"], ["X", "y", "treatment", "w"]),
+        # DRIV: `assignment` lands fourth, after X, y and treatment.
+        (
+            ["X", "assignment", "treatment", "y", "p", "pZ"],
+            ["X", "y", "treatment", "assignment", "p", "pZ"],
+        ),
+        # A signature with neither is left completely alone.
+        (["X", "sample_weight"], ["X", "sample_weight"]),
+    ],
+)
+def test_v1_order(current, expected):
+    """X first, then y, then treatment; everything else keeps its relative order."""
+    assert v1_order(current) == expected
+
+
+# --- coverage across the whole package --------------------------------------
+
+
+def _shimmable_methods():
+    """Yield (qualname, method) for every learner method taking treatment and y.
+
+    Walks the installed package so a newly added learner is picked up
+    automatically. Modules whose optional backend (tf/torch/jax) is missing are
+    skipped, so this covers them only in the backend CI lanes.
+    """
+    import causalml
+
+    for info in pkgutil.walk_packages(causalml.__path__, prefix="causalml."):
+        try:
+            module = importlib.import_module(info.name)
+        except Exception:  # optional backend absent, or an import-time failure
+            continue
+        for _, cls in inspect.getmembers(module, inspect.isclass):
+            if not cls.__module__.startswith("causalml."):
+                continue
+            for name in SHIMMED_METHODS:
+                method = cls.__dict__.get(name)
+                if not callable(method) or getattr(
+                    method, "__isabstractmethod__", False
+                ):
+                    continue
+                params = _positional_params(getattr(method, "__wrapped__", method))
+                if "treatment" in params and "y" in params:
+                    yield f"{cls.__module__}.{cls.__qualname__}.{name}", method
+
+
+def test_every_treatment_y_method_is_shimmed():
+    """Completeness guard for the deprecation window.
+
+    A deprecation window is one-shot: anything whose positional order changes at
+    v1.0 must warn in this release or it needs a second cycle. This fails if a
+    learner method takes ``treatment`` and ``y`` positionally without the shim.
+    """
+    unshimmed = [
+        qualname
+        for qualname, method in _shimmable_methods()
+        if not getattr(method, "_arg_order_shimmed", False)
+    ]
+    assert unshimmed == [], f"missing #854 shim on: {unshimmed}"
+
+
+def test_shim_covers_more_than_the_meta_learners():
+    """Guards against the hook silently regressing to BaseLearner-only scope."""
+    modules = {qualname.rsplit(".", 2)[0] for qualname, _ in _shimmable_methods()}
+    assert any("inference.tree" in m for m in modules), modules
+    assert any("inference.iv" in m for m in modules), modules
+
+
+# --- behaviour outside the meta-learner family ------------------------------
+
+
+def test_causal_tree_fit_warns_positionally(generate_regression_data):
+    """The trees inherit sklearn estimator bases, so they need the flip too."""
+    y, X, treatment, _, _, _ = generate_regression_data()
+    with pytest.warns(FutureWarning, match="argument order"):
+        CausalTreeRegressor().fit(X, treatment, y)
+
+
+def test_causal_tree_fit_keyword_is_silent(generate_regression_data):
+    y, X, treatment, _, _, _ = generate_regression_data()
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        CausalTreeRegressor().fit(X=X, y=y, treatment=treatment)
+    assert _order_warnings(record) == []
+
+
+def test_iv_regressor_fit_warns_positionally():
+    """IVRegressor's message must name its own signature, instrument included."""
+    rng = np.random.RandomState(RANDOM_SEED)
+    n = 200
+    w = rng.binomial(1, 0.5, n).astype(float)
+    treatment = w * rng.binomial(1, 0.8, n)
+    X = rng.normal(size=(n, 2))
+    y = treatment + X[:, 0] + rng.normal(size=n)
+    with pytest.warns(FutureWarning, match=r"becomes \(X, y, treatment, w\)"):
+        IVRegressor().fit(X, treatment, y, w)
+
+
+def test_driv_message_puts_assignment_fourth():
+    """Pins the v1.0 order chosen for DRIV: (X, y, treatment, assignment, ...)."""
+    params = _positional_params(BaseDRIVRegressor.fit.__wrapped__)
+    assert v1_order(params)[:4] == ["X", "y", "treatment", "assignment"]
+
+
+# --- predict joins the deprecation ------------------------------------------
+
+
+@pytest.mark.parametrize("learner_cls", PREDICT_REORDERED)
+def test_predict_warns_positionally(generate_regression_data, learner_cls):
+    """After the flip fit and predict must agree, so predict warns too."""
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = learner_cls(learner=LinearRegression())
+    learner.fit(X=X, y=y, treatment=treatment)
+    with pytest.warns(FutureWarning, match="argument order"):
+        learner.predict(X, treatment, y)
+
+
+@pytest.mark.parametrize("learner_cls", PREDICT_REORDERED)
+def test_predict_keyword_is_silent(generate_regression_data, learner_cls):
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = learner_cls(learner=LinearRegression())
+    learner.fit(X=X, y=y, treatment=treatment)
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        learner.predict(X, treatment=treatment, y=y)
+    assert _order_warnings(record) == []
+
+
+def test_fit_predict_warns_once_across_fit_and_predict(generate_regression_data):
+    """fit_predict delegates to both fit and predict; still one warning."""
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = BaseTRegressor(learner=LinearRegression())
+    with warnings.catch_warnings(record=True) as record:
+        warnings.simplefilter("always")
+        learner.fit_predict(X, treatment, y)
+    assert len(_order_warnings(record)) == 1
+
+
+# --- classifier fit overrides are wrapped independently ---------------------
+
+
+@pytest.mark.parametrize(
+    "learner_cls", [BaseSClassifier, BaseTClassifier, BaseXClassifier, BaseRClassifier]
+)
+def test_classifier_fit_is_shimmed(learner_cls):
+    """The classifier variants define their own fit, so each is wrapped separately."""
+    assert getattr(learner_cls.fit, "_arg_order_shimmed", False)
+    assert _positional_params(learner_cls.fit)[:3] == ["X", "treatment", "y"]

From 26678430647dd6fa58a39571466f141f0cbbf3e5 Mon Sep 17 00:00:00 2001
From: Jeong-Yoon Lee <jeongyoon.lee1@gmail.com>
Date: Fri, 31 Jul 2026 12:13:27 +0900
Subject: [PATCH 3/3] Close the #854 deprecation window: shim by signature, not
 by name (#981)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The shim keyed off four method names (fit/fit_predict/estimate_ate/predict).
A scan for public methods taking `treatment` before `y` found eleven more
that got no warning at all: four `bootstrap` variants,
`fit_bootstrap_ensemble`, `bootstrap_pool`, `prune`, `fill`, and the three
`Sensitivity.get_*` helpers. A deprecation window is one-shot, so a method
missing from it needs a second cycle of its own.

Replace the name allowlist with the property we actually care about: shim
any method whose positional order changes at v1.0. That fixes the class of
bug rather than the eleven instances, and the completeness test now scans
every public method instead of the same four names it was blind to.

Resolves both open questions on the epic:

- `UpliftTreeClassifier.fit`'s validation triple is reordered in place, so
  it stays consistent with the main one: (X, y, treatment, X_val, y_val,
  treatment_val, ...) rather than leaving X_val, treatment_val, y_val.
- The `Sensitivity` helpers are in scope. Their (X, p, treatment, y) is a
  third argument shape and becomes (X, y, treatment, p).

Two fixes the work required:

- `timeit` did not use functools.wraps, so `CausalTreeRegressor.bootstrap_pool`
  reported `(*args, **kw)` to `inspect.signature` and `"wrapped"` as its
  `__name__` — which is exactly why a signature-based shim could not see it.
  This also repairs introspection for anything sklearn does with that method.

- `v1_order` now reorders only when both `y` and `treatment` are present.
  Hoisting `y` on its own reached the vendored sklearn tree builders
  (`build(tree, X, y, ...)`, where it would land ahead of `X`) and
  `SemiSynthDataGenerator.fit(X, w, y, ...)`. Both are pinned as tests.

Tests 43 -> 52. Full suite: 418 passed, 10 skipped.

Closes #981. Part of #980.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 causalml/inference/_arg_order.py | 140 ++++++++++++++++++++++---------
 causalml/inference/tree/utils.py |   2 +
 causalml/metrics/sensitivity.py  |   3 +
 tests/test_fit_arg_order.py      | 101 ++++++++++++++++++----
 4 files changed, 192 insertions(+), 54 deletions(-)

diff --git a/causalml/inference/_arg_order.py b/causalml/inference/_arg_order.py
index f5ee5694..476da69b 100644
--- a/causalml/inference/_arg_order.py
+++ b/causalml/inference/_arg_order.py
@@ -3,15 +3,29 @@
 CausalML learners historically take ``fit(X, treatment, y, ...)``, which puts
 ``y`` third and breaks ``sklearn.pipeline.Pipeline`` (it calls the final
 estimator's ``fit(X, y)`` positionally). In v1.0 the positional order becomes
-``(X, y, treatment, ...)``; every other parameter keeps its relative position,
-so ``BaseDRIVLearner.fit(X, assignment, treatment, y, ...)`` becomes
-``fit(X, y, treatment, assignment, ...)``.
+``(X, y, treatment, ...)``:
+
+* ``y`` and ``treatment`` move to the front, immediately after ``X``, and every
+  other parameter keeps its relative position — so
+  ``BaseDRIVLearner.fit(X, assignment, treatment, y, ...)`` becomes
+  ``fit(X, y, treatment, assignment, ...)``.
+* A suffixed pair such as ``treatment_val``/``y_val`` is reordered in place, so
+  ``UpliftTreeClassifier.fit(X, treatment, y, X_val, treatment_val, y_val, ...)``
+  becomes ``fit(X, y, treatment, X_val, y_val, treatment_val, ...)`` rather than
+  leaving the validation triple inconsistent with the main one.
 
 This is a two-step deprecation. Step one (this module): the positional order is
 UNCHANGED, so no existing call silently breaks, but passing ``treatment`` or
 ``y`` positionally emits a ``FutureWarning`` steering callers to keyword
 arguments, which are order-independent and therefore safe across the flip.
 Step two (v1.0): reorder the signatures and delete this module.
+
+Coverage is decided by **signature, not by method name**. Any method whose
+positional order would change is shimmed — ``fit`` and ``predict``, but equally
+``bootstrap``, ``fit_bootstrap_ensemble``, ``prune`` and the ``Sensitivity``
+helpers. An earlier name-based allowlist silently missed eleven public methods
+(#981); the deprecation window is one-shot, so a method that is not shimmed here
+would need a second deprecation cycle of its own.
 """
 
 import functools
@@ -21,17 +35,6 @@
 #: Parameters whose positional slot changes in v1.0, in their v1.0 order.
 REORDERED_PARAMS = ("y", "treatment")
 
-#: Methods wrapped automatically for ``SerializableLearner`` subclasses.
-SHIMMED_METHODS = ("fit", "fit_predict", "estimate_ate", "predict")
-
-_MSG = (
-    "Passing `treatment` and/or `y` to {name}() by position is deprecated and "
-    "will change in causalml v1.0: the positional argument order becomes "
-    "({new_order}) for scikit-learn Pipeline compatibility "
-    "(see https://github.com/uber/causalml/issues/854). To be safe across the "
-    "change, pass them as keyword arguments, e.g. {example}."
-)
-
 
 def _positional_params(method):
     """Return the names of ``method``'s positional parameters, minus ``self``."""
@@ -45,23 +48,71 @@ def _positional_params(method):
     return names[1:] if names[:1] == ["self"] else names
 
 
-def v1_order(params):
-    """Reorder ``params`` to the v1.0 convention: X, then y, treatment, then rest.
+def _order_suffixed_pairs(names):
+    """Put ``y`` before ``treatment`` within each suffixed pair, in place.
 
-    Parameters other than ``y``/``treatment`` keep their relative order, so a
-    signature's extra arguments (``assignment``, ``p``, ``sample_weight``, ...)
-    simply shift right.
+    Only pairs sharing a suffix are touched (``treatment_val``/``y_val``), and
+    they keep the two slots they already occupy — unlike the bare
+    ``y``/``treatment`` pair, which is hoisted to the front by :func:`v1_order`.
+    """
+    ordered = list(names)
+    suffixes = {
+        name[len("treatment") :]
+        for name in ordered
+        if name.startswith("treatment") and name != "treatment"
+    }
+    for suffix in suffixes:
+        pair = ("y" + suffix, "treatment" + suffix)
+        if all(name in ordered for name in pair):
+            first, second = sorted(ordered.index(name) for name in pair)
+            ordered[first], ordered[second] = pair
+    return ordered
+
+
+def v1_order(params):
+    """Reorder ``params`` to the v1.0 convention.
+
+    The feature matrix stays first, then ``y`` and ``treatment``, then
+    everything else in its original relative order with any suffixed
+    ``treatment``/``y`` pair swapped so ``y`` leads.
+
+    ``params`` is returned unchanged unless **both** ``y`` and ``treatment``
+    appear after the first position. This deprecation is about that pair; a
+    method taking only one of them is not part of it, which keeps the rule off
+    the vendored sklearn tree builders (``build(tree, X, y, ...)``, where
+    hoisting ``y`` would put it ahead of ``X``) and off dataset helpers that
+    name their treatment vector something else.
     """
     head, tail = params[:1], params[1:]
-    moved = [name for name in REORDERED_PARAMS if name in tail]
-    return head + moved + [name for name in tail if name not in REORDERED_PARAMS]
+    if not all(name in tail for name in REORDERED_PARAMS):
+        return list(params)
+    moved = list(REORDERED_PARAMS)
+    rest = [name for name in tail if name not in REORDERED_PARAMS]
+    return head + moved + _order_suffixed_pairs(rest)
+
+
+def _may_reorder(method):
+    """Cheap pre-filter: could this method's positional order change at all?
+
+    Reads the code object directly so the common case costs no
+    :func:`inspect.signature` call — ``shim_arg_order`` scans every method of
+    every learner class at import time. Unwraps first, because
+    :func:`functools.wraps` copies ``__wrapped__`` but not ``__code__``, so a
+    decorated method (e.g. ``@timeit``) would otherwise look argument-less.
+    """
+    code = getattr(inspect.unwrap(method), "__code__", None)
+    if code is None:
+        return False
+    args = set(code.co_varnames[: code.co_argcount])
+    return bool(args & set(REORDERED_PARAMS))
 
 
 def deprecate_positional_treatment_y(method):
     """Warn when ``treatment``/``y`` are passed positionally to ``method``.
 
-    ``method`` is returned unchanged when neither parameter is positional, so
-    this is safe to apply blanket-style from ``__init_subclass__``.
+    ``method`` is returned unchanged when its positional order is already the
+    v1.0 order (including when it takes neither parameter), so this is safe to
+    apply blanket-style from ``__init_subclass__``.
 
     Only the outermost shimmed call on a given instance warns, so internal
     delegation (``fit_predict`` -> ``fit``/``predict``, a subclass ``fit`` ->
@@ -71,12 +122,19 @@ def deprecate_positional_treatment_y(method):
     scikit-learn metadata routing are unaffected.
     """
     params = _positional_params(method)
-    if not any(name in REORDERED_PARAMS for name in params):
+    target = v1_order(params)
+    if target == params:
         return method
 
-    message = _MSG.format(
+    message = (
+        "Passing `treatment` and/or `y` to {name}() by position is deprecated "
+        "and will change in causalml v1.0: the positional argument order "
+        "becomes ({new_order}) for scikit-learn Pipeline compatibility "
+        "(see https://github.com/uber/causalml/issues/854). To be safe across "
+        "the change, pass them as keyword arguments, e.g. {example}."
+    ).format(
         name=method.__name__,
-        new_order=", ".join(v1_order(params)),
+        new_order=", ".join(target),
         example="{}({}, {})".format(
             method.__name__,
             params[0],
@@ -103,22 +161,28 @@ def wrapper(self, *args, **kwargs):
 
 
 def shim_arg_order(cls):
-    """Wrap ``cls``'s own :data:`SHIMMED_METHODS` with the arg-order shim.
+    """Shim every method of ``cls`` whose positional order changes in v1.0.
 
-    Only methods defined directly on ``cls`` are wrapped (inherited ones are
-    already wrapped on the parent), abstract methods are left alone, and
-    already-wrapped methods are skipped, so nothing is double-wrapped.
+    Membership is decided by signature rather than by a list of method names —
+    a name list is what let ``bootstrap``, ``prune`` and the ``Sensitivity``
+    helpers slip through the window (#981). Only methods defined directly on
+    ``cls`` are considered (inherited ones are already wrapped on the parent);
+    dunders, ``classmethod``/``staticmethod`` descriptors, abstract methods and
+    already-wrapped methods are skipped.
 
     ``SerializableLearner`` applies this automatically via ``__init_subclass__``.
-    Learners outside that hierarchy — the TF/Torch/JAX estimators,
-    ``PolicyLearner``, ``TMLELearner`` — use it as a class decorator.
+    Classes outside that hierarchy — the TF/Torch/JAX estimators,
+    ``PolicyLearner``, ``TMLELearner``, ``Sensitivity`` — use it as a decorator.
     """
-    for name in SHIMMED_METHODS:
-        method = cls.__dict__.get(name)
+    for name, method in list(vars(cls).items()):
         if (
-            callable(method)
-            and not getattr(method, "_arg_order_shimmed", False)
-            and not getattr(method, "__isabstractmethod__", False)
+            name.startswith("_")
+            or isinstance(method, (classmethod, staticmethod))
+            or not callable(method)
+            or getattr(method, "__isabstractmethod__", False)
+            or getattr(method, "_arg_order_shimmed", False)
+            or not _may_reorder(method)
         ):
-            setattr(cls, name, deprecate_positional_treatment_y(method))
+            continue
+        setattr(cls, name, deprecate_positional_treatment_y(method))
     return cls
diff --git a/causalml/inference/tree/utils.py b/causalml/inference/tree/utils.py
index fbd22efe..0baa09d8 100644
--- a/causalml/inference/tree/utils.py
+++ b/causalml/inference/tree/utils.py
@@ -2,6 +2,7 @@
 Utility functions for uplift trees.
 """
 
+import functools
 import time
 from typing import Callable
 
@@ -342,6 +343,7 @@ def timeit(exclude_kwargs: tuple = ()) -> Callable:
     """
 
     def wrapper(f: Callable):
+        @functools.wraps(f)
         def wrapped(*args, **kw):
             ts = time.time()
             result = f(*args, **kw)
diff --git a/causalml/metrics/sensitivity.py b/causalml/metrics/sensitivity.py
index c0ac4206..2ef0830e 100644
--- a/causalml/metrics/sensitivity.py
+++ b/causalml/metrics/sensitivity.py
@@ -4,6 +4,8 @@
 import matplotlib.pyplot as plt
 from importlib import import_module
 
+from causalml.inference._arg_order import shim_arg_order
+
 logger = logging.getLogger("sensitivity")
 
 SUMMARY_COLS = ["Method", "ATE", "New ATE", "New ATE LB", "New ATE UB"]
@@ -95,6 +97,7 @@ def msm_propensity_bounds(p, gamma):
     return p_lower, p_upper
 
 
+@shim_arg_order
 class Sensitivity:
     """A Sensitivity Check class to support Placebo Treatment, Irrelevant Additional Confounder
     and Subset validation refutation methods to verify causal inference.
diff --git a/tests/test_fit_arg_order.py b/tests/test_fit_arg_order.py
index 00d918d2..c7ea4990 100644
--- a/tests/test_fit_arg_order.py
+++ b/tests/test_fit_arg_order.py
@@ -24,11 +24,7 @@
 import pytest
 from sklearn.linear_model import LinearRegression
 
-from causalml.inference._arg_order import (
-    SHIMMED_METHODS,
-    _positional_params,
-    v1_order,
-)
+from causalml.inference._arg_order import _positional_params, v1_order
 from causalml.inference.iv import BaseDRIVRegressor, IVRegressor
 from causalml.inference.meta import (
     BaseSRegressor,
@@ -42,7 +38,8 @@
     BaseXClassifier,
     BaseRClassifier,
 )
-from causalml.inference.tree import CausalTreeRegressor
+from causalml.inference.tree import CausalTreeRegressor, UpliftTreeClassifier
+from causalml.metrics.sensitivity import Sensitivity
 
 from .const import RANDOM_SEED
 
@@ -168,8 +165,21 @@ def test_fit_signature_is_preserved(learner_cls):
             ["X", "assignment", "treatment", "y", "p", "pZ"],
             ["X", "y", "treatment", "assignment", "p", "pZ"],
         ),
+        # UpliftTreeClassifier: the validation triple is reordered in place too,
+        # so it stays consistent with the main one (#980 open question 1).
+        (
+            ["X", "treatment", "y", "X_val", "treatment_val", "y_val", "sample_weight"],
+            ["X", "y", "treatment", "X_val", "y_val", "treatment_val", "sample_weight"],
+        ),
+        # Sensitivity helpers are in scope, and `p` shifts right (#980 q2).
+        (["X", "p", "treatment", "y"], ["X", "y", "treatment", "p"]),
         # A signature with neither is left completely alone.
         (["X", "sample_weight"], ["X", "sample_weight"]),
+        # Only `y`, no `treatment`: not this deprecation, so untouched. Keeps the
+        # rule off the vendored sklearn tree builders, where hoisting `y` would
+        # put it ahead of `X`.
+        (["tree", "X", "y", "sample_weight"], ["tree", "X", "y", "sample_weight"]),
+        (["X", "w", "y", "initial_taus"], ["X", "w", "y", "initial_taus"]),
     ],
 )
 def test_v1_order(current, expected):
@@ -181,7 +191,13 @@ def test_v1_order(current, expected):
 
 
 def _shimmable_methods():
-    """Yield (qualname, method) for every learner method taking treatment and y.
+    """Yield (qualname, method) for every public method whose order changes at v1.0.
+
+    Scans by **signature**, not by method name. An earlier name-based allowlist
+    silently missed eleven public methods (#981) — `bootstrap`,
+    `fit_bootstrap_ensemble`, `bootstrap_pool`, `prune`, `fill` and the three
+    `Sensitivity.get_*` helpers — so this guard must not be narrowed back to a
+    fixed set of names.
 
     Walks the installed package so a newly added learner is picked up
     automatically. Modules whose optional backend (tf/torch/jax) is missing are
@@ -197,23 +213,28 @@ def _shimmable_methods():
         for _, cls in inspect.getmembers(module, inspect.isclass):
             if not cls.__module__.startswith("causalml."):
                 continue
-            for name in SHIMMED_METHODS:
-                method = cls.__dict__.get(name)
-                if not callable(method) or getattr(
-                    method, "__isabstractmethod__", False
+            for name, method in vars(cls).items():
+                if (
+                    name.startswith("_")
+                    or isinstance(method, (classmethod, staticmethod))
+                    or not callable(method)
+                    or getattr(method, "__isabstractmethod__", False)
                 ):
                     continue
-                params = _positional_params(getattr(method, "__wrapped__", method))
-                if "treatment" in params and "y" in params:
+                try:
+                    params = _positional_params(getattr(method, "__wrapped__", method))
+                except (TypeError, ValueError):  # not introspectable
+                    continue
+                if v1_order(params) != params:
                     yield f"{cls.__module__}.{cls.__qualname__}.{name}", method
 
 
-def test_every_treatment_y_method_is_shimmed():
+def test_every_reordered_method_is_shimmed():
     """Completeness guard for the deprecation window.
 
     A deprecation window is one-shot: anything whose positional order changes at
-    v1.0 must warn in this release or it needs a second cycle. This fails if a
-    learner method takes ``treatment`` and ``y`` positionally without the shim.
+    v1.0 must warn in this release or it needs a second cycle. This fails if any
+    public method's order would change without the shim.
     """
     unshimmed = [
         qualname
@@ -223,6 +244,12 @@ def test_every_treatment_y_method_is_shimmed():
     assert unshimmed == [], f"missing #854 shim on: {unshimmed}"
 
 
+def test_shim_reaches_past_the_fit_family():
+    """#981: the shim must not regress to a fixed list of method names."""
+    names = {qualname.rsplit(".", 1)[1] for qualname, _ in _shimmable_methods()}
+    assert {"bootstrap", "fit_bootstrap_ensemble", "get_prediction"} <= names, names
+
+
 def test_shim_covers_more_than_the_meta_learners():
     """Guards against the hook silently regressing to BaseLearner-only scope."""
     modules = {qualname.rsplit(".", 2)[0] for qualname, _ in _shimmable_methods()}
@@ -310,3 +337,45 @@ def test_classifier_fit_is_shimmed(learner_cls):
     """The classifier variants define their own fit, so each is wrapped separately."""
     assert getattr(learner_cls.fit, "_arg_order_shimmed", False)
     assert _positional_params(learner_cls.fit)[:3] == ["X", "treatment", "y"]
+
+
+# --- #981: coverage beyond the fit family -----------------------------------
+
+
+def test_bootstrap_warns_positionally(generate_regression_data):
+    """`bootstrap` is public API on the most-used learners and changes at v1.0."""
+    y, X, treatment, _, _, _ = generate_regression_data()
+    learner = BaseTRegressor(learner=LinearRegression())
+    learner.fit(X=X, y=y, treatment=treatment)
+    with pytest.warns(FutureWarning, match="argument order"):
+        learner.bootstrap(X, treatment, y, None, 200)
+
+
+def test_sensitivity_get_prediction_message_moves_p_right():
+    """Sensitivity takes (X, p, treatment, y) — a third shape (#980 q2)."""
+    params = _positional_params(Sensitivity.get_prediction.__wrapped__)
+    assert params == ["X", "p", "treatment", "y"]
+    assert v1_order(params) == ["X", "y", "treatment", "p"]
+
+
+def test_uplift_tree_validation_triple_is_reordered():
+    """#980 q1: X_val/y_val/treatment_val stays consistent with the main triple."""
+    params = _positional_params(UpliftTreeClassifier.fit.__wrapped__)
+    assert v1_order(params) == [
+        "X",
+        "y",
+        "treatment",
+        "X_val",
+        "y_val",
+        "treatment_val",
+        "sample_weight",
+        "check_input",
+    ]
+
+
+def test_timeit_preserves_signature():
+    """`bootstrap_pool` is behind @timeit; without functools.wraps its signature
+    reads (*args, **kw), which hid it from the signature-based shim."""
+    assert CausalTreeRegressor.bootstrap_pool.__name__ == "bootstrap_pool"
+    params = _positional_params(CausalTreeRegressor.bootstrap_pool.__wrapped__)
+    assert params[:3] == ["X", "treatment", "y"]