From c106ffa1f280803c9583fb96f26f4369721afaaa Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Sat, 23 May 2026 12:56:34 +0530
Subject: [PATCH 01/15] add Polars DataFrame/Series/LazyFrame support

---
 causalml/inference/meta/utils.py | 103 +++++++++-
 causalml/propensity.py           |  21 ++-
 tests/test_polars_support.py     | 314 +++++++++++++++++++++++++++++++
 3 files changed, 425 insertions(+), 13 deletions(-)
 create mode 100644 tests/test_polars_support.py

diff --git a/causalml/inference/meta/utils.py b/causalml/inference/meta/utils.py
index 157eeaf6..be0ceca7 100644
--- a/causalml/inference/meta/utils.py
+++ b/causalml/inference/meta/utils.py
@@ -4,9 +4,81 @@
 from packaging import version
 from xgboost import __version__ as xgboost_version
 
+# Optional Polars import
+try:
+    import polars as pl
+
+    _POLARS_AVAILABLE = True
+except ImportError:
+    pl = None
+    _POLARS_AVAILABLE = False
+
+
+def _is_polars_dataframe(obj):
+    """Return True if *obj* is a polars DataFrame or LazyFrame."""
+    if not _POLARS_AVAILABLE:
+        return False
+    return isinstance(obj, (pl.DataFrame, pl.LazyFrame))
+
+
+def _is_polars_series(obj):
+    """Return True if *obj* is a polars Series."""
+    if not _POLARS_AVAILABLE:
+        return False
+    return isinstance(obj, pl.Series)
+
+
+def _polars_to_numpy(obj):
+    """Convert a polars DataFrame, LazyFrame, or Series to a NumPy array.
+
+    - ``pl.LazyFrame`` is collected first (implicit ``.collect()``).
+    - A single-column ``pl.DataFrame`` is squeezed to a 1-D array to match
+      the behaviour of ``pd.Series.to_numpy()``.
+    - A multi-column ``pl.DataFrame`` is returned as a 2-D array.
+    """
+    if isinstance(obj, pl.LazyFrame):
+        obj = obj.collect()
+
+    if isinstance(obj, pl.DataFrame):
+        arr = obj.to_numpy()
+        # Squeeze single-column frames so downstream code gets a 1-D vector
+        if arr.shape[1] == 1:
+            return arr.ravel()
+        return arr
+
+    if isinstance(obj, pl.Series):
+        return obj.to_numpy()
+
+    raise TypeError(f"Expected a polars DataFrame/LazyFrame/Series, got {type(obj)}")
+
+
+# ---------------------------------------------------------------------------
+# Public helpers
+# ---------------------------------------------------------------------------
+
 
 def convert_pd_to_np(*args):
-    output = [obj.to_numpy() if hasattr(obj, "to_numpy") else obj for obj in args]
+    """Convert pandas or polars objects to NumPy arrays.
+
+    Accepts any mix of:
+    * ``pd.DataFrame`` / ``pd.Series``   → ``.to_numpy()``
+    * ``pl.DataFrame`` / ``pl.LazyFrame`` / ``pl.Series``  → converted via
+      :func:`_polars_to_numpy`
+    * Any other type (e.g. ``np.ndarray``, ``None``) → returned unchanged.
+    """
+
+    def _convert(obj):
+        if obj is None:
+            return obj
+        if _POLARS_AVAILABLE and isinstance(
+            obj, (pl.DataFrame, pl.LazyFrame, pl.Series)
+        ):
+            return _polars_to_numpy(obj)
+        if hasattr(obj, "to_numpy"):
+            return obj.to_numpy()
+        return obj
+
+    output = [_convert(obj) for obj in args]
     return output if len(output) > 1 else output[0]
 
 
@@ -21,21 +93,34 @@ def check_treatment_vector(treatment, control_name=None):
 
 def check_p_conditions(p, t_groups):
     eps = np.finfo(float).eps
+
+    # Build the allowed types tuple dynamically so it works whether or not
+    # polars is installed.
+    _allowed = [np.ndarray, pd.Series]
+    if _POLARS_AVAILABLE:
+        _allowed.append(pl.Series)
+    _allowed_tuple = tuple(_allowed)
+
     assert isinstance(
-        p, (np.ndarray, pd.Series, dict)
-    ), "p must be an np.ndarray, pd.Series, or dict type"
-    if isinstance(p, (np.ndarray, pd.Series)):
+        p, (*_allowed_tuple, dict)
+    ), "p must be an np.ndarray, pd.Series, pl.Series (if polars is installed), or dict type"
+
+    if isinstance(p, _allowed_tuple):
+        # Normalise to numpy for the value checks below
+        p_np = p.to_numpy() if hasattr(p, "to_numpy") else np.asarray(p)
         assert (
             t_groups.shape[0] == 1
-        ), "If p is passed as an np.ndarray, there must be only 1 unique non-control group in the treatment vector."
-        assert (0 + eps < p).all() and (
-            p < 1 - eps
+        ), "If p is passed as an array/Series, there must be only 1 unique non-control group in the treatment vector."
+        assert (0 + eps < p_np).all() and (
+            p_np < 1 - eps
         ).all(), "The values of p should lie within the (0, 1) interval."
 
     if isinstance(p, dict):
         for t_name in t_groups:
-            assert (0 + eps < p[t_name]).all() and (
-                p[t_name] < 1 - eps
+            p_val = p[t_name]
+            p_np = p_val.to_numpy() if hasattr(p_val, "to_numpy") else np.asarray(p_val)
+            assert (0 + eps < p_np).all() and (
+                p_np < 1 - eps
             ).all(), "The values of p should lie within the (0, 1) interval."
 
 
diff --git a/causalml/propensity.py b/causalml/propensity.py
index 4e12dcb9..ec6aa981 100644
--- a/causalml/propensity.py
+++ b/causalml/propensity.py
@@ -7,6 +7,8 @@
 from sklearn.isotonic import IsotonicRegression
 import xgboost as xgb
 
+from causalml.inference.meta.utils import convert_pd_to_np
+
 logger = logging.getLogger("causalml")
 
 
@@ -41,6 +43,7 @@ def fit(self, X, y):
             X (numpy.ndarray): a feature matrix
             y (numpy.ndarray): a binary target vector
         """
+        X, y = convert_pd_to_np(X, y)
         self.model.fit(X, y)
         if self.calibrate:
             # Fit a calibrator to the propensity scores with IsotonicRegression.
@@ -62,6 +65,7 @@ def predict(self, X):
         Returns:
             (numpy.ndarray): Propensity scores between 0 and 1.
         """
+        X = convert_pd_to_np(X)
         p = self.model.predict_proba(X)[:, 1]
         if self.calibrate:
             p = self.calibrator.transform(p)
@@ -162,6 +166,7 @@ def fit(self, X, y, stop_val_size=0.2):
             X (numpy.ndarray): a feature matrix
             y (numpy.ndarray): a binary target vector
         """
+        X, y = convert_pd_to_np(X, y)
 
         if self.early_stop:
             X_train, X_val, y_train, y_val = train_test_split(
@@ -197,21 +202,29 @@ def compute_propensity_score(
 
     Args:
         X (np.matrix): features for training
-        treatment (np.array or pd.Series): a treatment vector for training
+        treatment (np.array or pd.Series or pl.Series): a treatment vector for training
         p_model (model object, optional): a binary classifier with either a predict_proba or predict method
         X_pred (np.matrix, optional): features for prediction
-        treatment_pred (np.array or pd.Series, optional): a treatment vector for prediciton
+        treatment_pred (np.array or pd.Series or pl.Series, optional): a treatment vector for prediction
         calibrate_p (bool, optional): whether calibrate the propensity score
-        clip_bounds (tuple, optional): lower and upper bounds for clipping propensity scores. Bounds should be implemented
-                    such that: 0 < lower < upper < 1, to avoid division by zero in BaseRLearner.fit_predict() step.
+        clip_bounds (tuple, optional): lower and upper bounds for clipping propensity scores.
 
     Returns:
         (tuple)
             - p (numpy.ndarray): propensity score
             - p_model (PropensityModel): either the original p_model or a trained ElasticNetPropensityModel
     """
+    # Normalise inputs to numpy so downstream sklearn models always see arrays.
+    X, treatment = convert_pd_to_np(X, treatment)
+
     if treatment_pred is None:
         treatment_pred = treatment.copy()
+    else:
+        treatment_pred = convert_pd_to_np(treatment_pred)
+
+    if X_pred is not None:
+        X_pred = convert_pd_to_np(X_pred)
+
     if p_model is None:
         p_model = ElasticNetPropensityModel(
             clip_bounds=clip_bounds, calibrate=calibrate_p
diff --git a/tests/test_polars_support.py b/tests/test_polars_support.py
new file mode 100644
index 00000000..1d038926
--- /dev/null
+++ b/tests/test_polars_support.py
@@ -0,0 +1,314 @@
+"""
+Tests for Polars DataFrame/Series support across CausalML meta-learners.
+
+Run with:
+    pytest tests/test_polars_support.py -v
+"""
+
+import pytest
+import numpy as np
+import pandas as pd
+
+# Skip the entire module if polars is not installed
+polars = pytest.importorskip("polars", reason="polars is not installed")
+import polars as pl
+
+from sklearn.linear_model import LinearRegression
+
+from causalml.inference.meta.tlearner import BaseTRegressor
+from causalml.inference.meta.slearner import BaseSRegressor
+from causalml.inference.meta.xlearner import BaseXRegressor
+from causalml.inference.meta.rlearner import BaseRRegressor
+from causalml.inference.meta.drlearner import BaseDRRegressor
+from causalml.inference.meta.utils import convert_pd_to_np, check_p_conditions
+
+# Fixtures
+
+N = 200
+N_FEATURES = 5
+RANDOM_STATE = 42
+
+
+@pytest.fixture(scope="module")
+def synthetic_data_numpy():
+    """Return (X, treatment, y) as NumPy arrays — the baseline."""
+    rng = np.random.default_rng(RANDOM_STATE)
+    X = rng.standard_normal((N, N_FEATURES))
+    treatment = rng.choice([0, 1], size=N)
+    y = X[:, 0] * treatment + rng.standard_normal(N) * 0.1
+    return X, treatment, y
+
+
+@pytest.fixture(scope="module")
+def synthetic_data_pandas(synthetic_data_numpy):
+    """Return (X, treatment, y) as pandas objects."""
+    X_np, t_np, y_np = synthetic_data_numpy
+    X = pd.DataFrame(X_np, columns=[f"f{i}" for i in range(N_FEATURES)])
+    treatment = pd.Series(t_np, name="treatment")
+    y = pd.Series(y_np, name="outcome")
+    return X, treatment, y
+
+
+@pytest.fixture(scope="module")
+def synthetic_data_polars(synthetic_data_numpy):
+    """Return (X, treatment, y) as polars objects."""
+    X_np, t_np, y_np = synthetic_data_numpy
+    X = pl.DataFrame({f"f{i}": X_np[:, i] for i in range(N_FEATURES)})
+    treatment = pl.Series("treatment", t_np)
+    y = pl.Series("outcome", y_np)
+    return X, treatment, y
+
+
+@pytest.fixture(scope="module")
+def synthetic_data_polars_lazy(synthetic_data_polars):
+    """Return X as a polars LazyFrame, treatment and y as Series."""
+    X, treatment, y = synthetic_data_polars
+    return X.lazy(), treatment, y
+
+
+# convert_pd_to_np unit tests
+
+
+class TestConvertPdToNp:
+    def test_numpy_passthrough(self, synthetic_data_numpy):
+        X, t, y = synthetic_data_numpy
+        X_out, t_out, y_out = convert_pd_to_np(X, t, y)
+        np.testing.assert_array_equal(X_out, X)
+        np.testing.assert_array_equal(t_out, t)
+
+    def test_pandas_conversion(self, synthetic_data_pandas):
+        X, t, y = synthetic_data_pandas
+        X_out, t_out, y_out = convert_pd_to_np(X, t, y)
+        assert isinstance(X_out, np.ndarray)
+        assert isinstance(t_out, np.ndarray)
+        assert isinstance(y_out, np.ndarray)
+
+    def test_polars_dataframe_conversion(self, synthetic_data_polars):
+        X, t, y = synthetic_data_polars
+        X_out, t_out, y_out = convert_pd_to_np(X, t, y)
+        assert isinstance(X_out, np.ndarray)
+        assert X_out.shape == (N, N_FEATURES)
+        assert isinstance(t_out, np.ndarray)
+        assert t_out.shape == (N,)
+        assert isinstance(y_out, np.ndarray)
+
+    def test_polars_lazyframe_conversion(self, synthetic_data_polars_lazy):
+        X_lazy, t, y = synthetic_data_polars_lazy
+        X_out = convert_pd_to_np(X_lazy)
+        assert isinstance(X_out, np.ndarray)
+        assert X_out.shape == (N, N_FEATURES)
+
+    def test_none_passthrough(self):
+        result = convert_pd_to_np(None)
+        assert result is None
+
+    def test_single_arg(self, synthetic_data_polars):
+        X, _, _ = synthetic_data_polars
+        X_out = convert_pd_to_np(X)
+        assert isinstance(X_out, np.ndarray)
+
+    def test_single_column_polars_df_is_1d(self):
+        """A single-column pl.DataFrame should be squeezed to 1-D."""
+        s = pl.DataFrame({"a": [1, 2, 3]})
+        out = convert_pd_to_np(s)
+        assert out.ndim == 1
+
+    def test_multi_column_polars_df_is_2d(self, synthetic_data_polars):
+        X, _, _ = synthetic_data_polars
+        out = convert_pd_to_np(X)
+        assert out.ndim == 2
+
+
+# check_p_conditions with polars Series
+
+
+class TestCheckPConditions:
+    def test_polars_series_accepted(self):
+        t_groups = np.array([1])
+        p = pl.Series("p", np.linspace(0.1, 0.9, N))
+        # Should not raise
+        check_p_conditions(p, t_groups)
+
+    def test_polars_series_out_of_bounds_raises(self):
+        t_groups = np.array([1])
+        p = pl.Series("p", np.linspace(0.0, 1.0, N))  # includes 0 and 1
+        with pytest.raises(AssertionError):
+            check_p_conditions(p, t_groups)
+
+
+# Helper: assert predictions are close between two input formats
+
+
+def _assert_te_close(te_ref, te_other, atol=1e-5):
+    """Treatment-effect arrays from two input formats should be identical."""
+    np.testing.assert_allclose(
+        te_ref,
+        te_other,
+        atol=atol,
+        err_msg="Treatment effects differ between input formats",
+    )
+
+
+# T-Learner
+
+
+class TestTLearnerPolars:
+    @pytest.fixture(autouse=True)
+    def _learner(self):
+        self.learner = BaseTRegressor(learner=LinearRegression())
+
+    def _fit_predict(self, X, treatment, y):
+        self.learner.fit(X, treatment, y)
+        return self.learner.predict(X)
+
+    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        # Re-init so models are fresh
+        self.learner = BaseTRegressor(learner=LinearRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_np, te_pl)
+
+    def test_polars_matches_pandas(self, synthetic_data_pandas, synthetic_data_polars):
+        te_pd = self._fit_predict(*synthetic_data_pandas)
+        self.learner = BaseTRegressor(learner=LinearRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_pd, te_pl)
+
+    def test_lazyframe_input(self, synthetic_data_numpy, synthetic_data_polars_lazy):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseTRegressor(learner=LinearRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars_lazy)
+        _assert_te_close(te_np, te_pl)
+
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
+        te = self._fit_predict(*synthetic_data_polars)
+        assert isinstance(te, np.ndarray)
+
+    def test_estimate_ate_polars(self, synthetic_data_polars):
+        X, treatment, y = synthetic_data_polars
+        ate, lb, ub = self.learner.estimate_ate(X, treatment, y)
+        assert isinstance(ate, np.ndarray)
+        assert lb[0] < ate[0] < ub[0]
+
+
+# S-Learner
+
+
+class TestSLearnerPolars:
+    @pytest.fixture(autouse=True)
+    def _learner(self):
+        self.learner = BaseSRegressor(learner=LinearRegression())
+
+    def _fit_predict(self, X, treatment, y):
+        self.learner.fit(X, treatment, y)
+        return self.learner.predict(X)
+
+    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseSRegressor(learner=LinearRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_np, te_pl)
+
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
+        te = self._fit_predict(*synthetic_data_polars)
+        assert isinstance(te, np.ndarray)
+
+    def test_estimate_ate_polars(self, synthetic_data_polars):
+        X, treatment, y = synthetic_data_polars
+        ate, lb, ub = self.learner.estimate_ate(X, treatment, y)
+        assert isinstance(ate, np.ndarray)
+
+
+# X-Learner
+
+
+class TestXLearnerPolars:
+    @pytest.fixture(autouse=True)
+    def _learner(self):
+        self.learner = BaseXRegressor(learner=LinearRegression())
+
+    def _fit_predict(self, X, treatment, y):
+        self.learner.fit(X, treatment, y)
+        return self.learner.predict(X)
+
+    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseXRegressor(learner=LinearRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_np, te_pl)
+
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
+        te = self._fit_predict(*synthetic_data_polars)
+        assert isinstance(te, np.ndarray)
+
+
+# R-Learner
+
+
+class TestRLearnerPolars:
+    @pytest.fixture(autouse=True)
+    def _learner(self):
+        self.learner = BaseRRegressor(learner=LinearRegression())
+
+    def _fit_predict(self, X, treatment, y):
+        self.learner.fit(X, treatment, y)
+        return self.learner.predict(X)
+
+    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseRRegressor(learner=LinearRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_np, te_pl)
+
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
+        te = self._fit_predict(*synthetic_data_polars)
+        assert isinstance(te, np.ndarray)
+
+
+# DR-Learner
+
+
+class TestDRLearnerPolars:
+    @pytest.fixture(autouse=True)
+    def _learner(self):
+        self.learner = BaseDRRegressor(learner=LinearRegression())
+
+    def _fit_predict(self, X, treatment, y):
+        self.learner.fit(X, treatment, y)
+        return self.learner.predict(X)
+
+    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseDRRegressor(learner=LinearRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_np, te_pl)
+
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
+        te = self._fit_predict(*synthetic_data_polars)
+        assert isinstance(te, np.ndarray)
+
+
+# Edge cases
+
+
+class TestEdgeCases:
+    def test_mixed_inputs_polars_x_numpy_treatment(
+        self, synthetic_data_numpy, synthetic_data_polars
+    ):
+        """X as polars DataFrame, treatment and y as numpy — should work fine."""
+        X_pl, _, _ = synthetic_data_polars
+        _, t_np, y_np = synthetic_data_numpy
+        learner = BaseTRegressor(learner=LinearRegression())
+        learner.fit(X_pl, t_np, y_np)
+        te = learner.predict(X_pl)
+        assert isinstance(te, np.ndarray)
+
+    def test_polars_predict_only(self, synthetic_data_numpy, synthetic_data_polars):
+        """Fit on numpy, predict on polars — predict must accept polars X."""
+        X_np, t_np, y_np = synthetic_data_numpy
+        X_pl, _, _ = synthetic_data_polars
+        learner = BaseTRegressor(learner=LinearRegression())
+        learner.fit(X_np, t_np, y_np)
+        te = learner.predict(X_pl)
+        assert isinstance(te, np.ndarray)
+        assert te.shape[0] == N

From c7b22a53f4c115b34b274a3547bc3f636d65cf18 Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Sat, 23 May 2026 13:05:08 +0530
Subject: [PATCH 02/15] improved test

---
 tests/test_polars_support.py | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/tests/test_polars_support.py b/tests/test_polars_support.py
index 1d038926..949cd91f 100644
--- a/tests/test_polars_support.py
+++ b/tests/test_polars_support.py
@@ -2,7 +2,7 @@
 Tests for Polars DataFrame/Series support across CausalML meta-learners.
 
 Run with:
-    pytest tests/test_polars_support.py -v
+    pytest tests/test_polars_support.py -v --noconftest
 """
 
 import pytest
@@ -163,7 +163,6 @@ def _fit_predict(self, X, treatment, y):
 
     def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
         te_np = self._fit_predict(*synthetic_data_numpy)
-        # Re-init so models are fresh
         self.learner = BaseTRegressor(learner=LinearRegression())
         te_pl = self._fit_predict(*synthetic_data_polars)
         _assert_te_close(te_np, te_pl)
@@ -215,8 +214,12 @@ def test_fit_predict_returns_numpy(self, synthetic_data_polars):
 
     def test_estimate_ate_polars(self, synthetic_data_polars):
         X, treatment, y = synthetic_data_polars
-        ate, lb, ub = self.learner.estimate_ate(X, treatment, y)
+        # BaseSLearner.estimate_ate returns only `ate` when return_ci=False (default)
+        ate = self.learner.estimate_ate(X, treatment, y)
         assert isinstance(ate, np.ndarray)
+        # With return_ci=True it returns (ate, lb, ub)
+        ate, lb, ub = self.learner.estimate_ate(X, treatment, y, return_ci=True)
+        assert lb[0] < ate[0] < ub[0]
 
 
 # X-Learner
@@ -248,15 +251,22 @@ def test_fit_predict_returns_numpy(self, synthetic_data_polars):
 class TestRLearnerPolars:
     @pytest.fixture(autouse=True)
     def _learner(self):
-        self.learner = BaseRRegressor(learner=LinearRegression())
+        # fixed random_state so KFold splits are identical across both runs
+        self.learner = BaseRRegressor(
+            learner=LinearRegression(), random_state=RANDOM_STATE
+        )
 
     def _fit_predict(self, X, treatment, y):
         self.learner.fit(X, treatment, y)
         return self.learner.predict(X)
 
     def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        # With a fixed random_state the KFold splits are deterministic,
+        # so numpy and polars inputs must produce identical results.
         te_np = self._fit_predict(*synthetic_data_numpy)
-        self.learner = BaseRRegressor(learner=LinearRegression())
+        self.learner = BaseRRegressor(
+            learner=LinearRegression(), random_state=RANDOM_STATE
+        )
         te_pl = self._fit_predict(*synthetic_data_polars)
         _assert_te_close(te_np, te_pl)
 
@@ -273,14 +283,16 @@ class TestDRLearnerPolars:
     def _learner(self):
         self.learner = BaseDRRegressor(learner=LinearRegression())
 
-    def _fit_predict(self, X, treatment, y):
-        self.learner.fit(X, treatment, y)
+    def _fit_predict(self, X, treatment, y, seed=None):
+        self.learner.fit(X, treatment, y, seed=seed)
         return self.learner.predict(X)
 
     def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
-        te_np = self._fit_predict(*synthetic_data_numpy)
+        # DR-Learner uses KFold with a seed parameter passed to fit(); fix it
+        # so both runs use the same splits.
+        te_np = self._fit_predict(*synthetic_data_numpy, seed=RANDOM_STATE)
         self.learner = BaseDRRegressor(learner=LinearRegression())
-        te_pl = self._fit_predict(*synthetic_data_polars)
+        te_pl = self._fit_predict(*synthetic_data_polars, seed=RANDOM_STATE)
         _assert_te_close(te_np, te_pl)
 
     def test_fit_predict_returns_numpy(self, synthetic_data_polars):

From 06d4a37f909b57bcc12e4f93cb80942fe43acb9c Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Wed, 3 Jun 2026 02:16:38 +0530
Subject: [PATCH 03/15] Polars support across all meta-learners

---
 causalml/inference/meta/drlearner.py | 364 ++++++++-----------------
 causalml/inference/meta/rlearner.py  | 337 +++++++----------------
 causalml/inference/meta/slearner.py  | 196 ++++----------
 causalml/inference/meta/tlearner.py  | 154 ++++++-----
 causalml/inference/meta/utils.py     | 240 ++++++++++-------
 causalml/inference/meta/xlearner.py  | 390 ++++++++++-----------------
 6 files changed, 627 insertions(+), 1054 deletions(-)

diff --git a/causalml/inference/meta/drlearner.py b/causalml/inference/meta/drlearner.py
index ea51b234..58640c4e 100644
--- a/causalml/inference/meta/drlearner.py
+++ b/causalml/inference/meta/drlearner.py
@@ -11,7 +11,9 @@
 from causalml.inference.meta.utils import (
     check_treatment_vector,
     check_p_conditions,
-    convert_pd_to_np,
+    filter_mask,
+    filter_index,
+    to_numpy,
 )
 from causalml.metrics import regression_metrics, classification_metrics
 from causalml.propensity import compute_propensity_score
@@ -20,12 +22,7 @@
 
 
 class BaseDRLearner(BaseLearner):
-    """A parent class for DR-learner regressor classes.
-
-    A DR-learner estimates treatment effects with machine learning models.
-
-    Details of DR-learner are available at `Kennedy (2020) <https://arxiv.org/abs/2004.14497>`_.
-    """
+    """A parent class for DR-learner regressor classes."""
 
     def __init__(
         self,
@@ -36,41 +33,30 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
-        """Initialize a DR-learner.
-
-        Args:
-            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
-                groups
-            control_outcome_learner (optional): a model to estimate outcomes in the control group
-            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
-            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         assert (learner is not None) or (
             (control_outcome_learner is not None)
             and (treatment_outcome_learner is not None)
             and (treatment_effect_learner is not None)
         )
 
-        if control_outcome_learner is None:
-            self.model_mu_c = deepcopy(learner)
-        else:
-            self.model_mu_c = control_outcome_learner
-
-        if treatment_outcome_learner is None:
-            self.model_mu_t = deepcopy(learner)
-        else:
-            self.model_mu_t = treatment_outcome_learner
-
-        if treatment_effect_learner is None:
-            self.model_tau = deepcopy(learner)
-        else:
-            self.model_tau = treatment_effect_learner
+        self.model_mu_c = (
+            deepcopy(learner)
+            if control_outcome_learner is None
+            else control_outcome_learner
+        )
+        self.model_mu_t = (
+            deepcopy(learner)
+            if treatment_outcome_learner is None
+            else treatment_outcome_learner
+        )
+        self.model_tau = (
+            deepcopy(learner)
+            if treatment_effect_learner is None
+            else treatment_effect_learner
+        )
 
         self.ate_alpha = ate_alpha
         self.control_name = control_name
-
         self.propensity = None
 
     def __repr__(self):
@@ -86,78 +72,63 @@ def __repr__(self):
         )
 
     def fit(self, X, treatment, y, p=None, seed=None):
-        """Fit the inference model.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            seed (int): random seed for cross-fitting
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
+        self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
 
-        # The estimator splits the data into 3 partitions for cross-fit on the propensity score estimation,
-        # the outcome regression, and the treatment regression on the doubly robust estimates. The use of
-        # the partitions is rotated so we do not lose on the sample size.
         cv = KFold(n_splits=3, shuffle=True, random_state=seed)
-        split_indices = [index for _, index in cv.split(y)]
+        split_indices = [index for _, index in cv.split(y_np)]
 
-        self.models_mu_c = [
-            deepcopy(self.model_mu_c),
-            deepcopy(self.model_mu_c),
-            deepcopy(self.model_mu_c),
-        ]
+        self.models_mu_c = [deepcopy(self.model_mu_c) for _ in range(3)]
         self.models_mu_t = {
-            group: [
-                deepcopy(self.model_mu_t),
-                deepcopy(self.model_mu_t),
-                deepcopy(self.model_mu_t),
-            ]
+            group: [deepcopy(self.model_mu_t) for _ in range(3)]
             for group in self.t_groups
         }
         self.models_tau = {
-            group: [
-                deepcopy(self.model_tau),
-                deepcopy(self.model_tau),
-                deepcopy(self.model_tau),
-            ]
+            group: [deepcopy(self.model_tau) for _ in range(3)]
             for group in self.t_groups
         }
+
         if p is None:
-            self.propensity = {group: np.zeros(y.shape[0]) for group in self.t_groups}
+            self.propensity = {
+                group: np.zeros(y_np.shape[0]) for group in self.t_groups
+            }
 
         for ifold in range(3):
             treatment_idx = split_indices[ifold]
             outcome_idx = split_indices[(ifold + 1) % 3]
             tau_idx = split_indices[(ifold + 2) % 3]
 
-            treatment_treat, treatment_out, treatment_tau = (
-                treatment[treatment_idx],
-                treatment[outcome_idx],
-                treatment[tau_idx],
-            )
-            y_out, y_tau = y[outcome_idx], y[tau_idx]
-            X_treat, X_out, X_tau = X[treatment_idx], X[outcome_idx], X[tau_idx]
+            treatment_treat = filter_index(treatment, treatment_idx)
+            treatment_out = filter_index(treatment, outcome_idx)
+            treatment_tau = filter_index(treatment, tau_idx)
+
+            treatment_treat_np = to_numpy(treatment_treat)
+            treatment_out_np = to_numpy(treatment_out)
+            treatment_tau_np = to_numpy(treatment_tau)
+
+            y_out = y_np[outcome_idx]
+            y_tau = y_np[tau_idx]
+
+            X_treat = filter_index(X, treatment_idx)
+            X_out = filter_index(X, outcome_idx)
+            X_tau = filter_index(X, tau_idx)
 
             if p is None:
                 logger.info("Generating propensity score")
                 cur_p = dict()
-
                 for group in self.t_groups:
-                    mask = (treatment_treat == group) | (
-                        treatment_treat == self.control_name
+                    mask = (treatment_treat_np == group) | (
+                        treatment_treat_np == self.control_name
                     )
-                    treatment_filt = treatment_treat[mask]
-                    X_filt = X_treat[mask]
-                    w_filt = (treatment_filt == group).astype(int)
-                    w = (treatment_tau == group).astype(int)
+                    treatment_filt = filter_mask(treatment_treat, mask)
+                    X_filt = filter_mask(X_treat, mask)
+                    w_filt = (to_numpy(treatment_filt) == group).astype(int)
+                    w = (treatment_tau_np == group).astype(int)
                     cur_p[group], _ = compute_propensity_score(
                         X=X_filt, treatment=w_filt, X_pred=X_tau, treatment_pred=w
                     )
@@ -165,29 +136,34 @@ def fit(self, X, treatment, y, p=None, seed=None):
             else:
                 cur_p = dict()
                 if isinstance(p, (np.ndarray, pd.Series)):
-                    cur_p = {self.t_groups[0]: convert_pd_to_np(p[tau_idx])}
+                    cur_p = {self.t_groups[0]: to_numpy(filter_index(p, tau_idx))}
                 else:
-                    cur_p = {g: prop[tau_idx] for g, prop in p.items()}
+                    cur_p = {
+                        g: to_numpy(filter_index(prop, tau_idx))
+                        for g, prop in p.items()
+                    }
                 check_p_conditions(cur_p, self.t_groups)
 
             logger.info("Generate outcome regressions")
             self.models_mu_c[ifold].fit(
-                X_out[treatment_out == self.control_name],
-                y_out[treatment_out == self.control_name],
+                filter_mask(X_out, treatment_out_np == self.control_name),
+                y_out[treatment_out_np == self.control_name],
             )
             for group in self.t_groups:
                 self.models_mu_t[group][ifold].fit(
-                    X_out[treatment_out == group], y_out[treatment_out == group]
+                    filter_mask(X_out, treatment_out_np == group),
+                    y_out[treatment_out_np == group],
                 )
 
             logger.info("Fit pseudo outcomes from the DR formula")
-
             for group in self.t_groups:
-                mask = (treatment_tau == group) | (treatment_tau == self.control_name)
-                treatment_filt = treatment_tau[mask]
-                X_filt = X_tau[mask]
+                mask = (treatment_tau_np == group) | (
+                    treatment_tau_np == self.control_name
+                )
+                treatment_filt_np = treatment_tau_np[mask]
+                X_filt = filter_mask(X_tau, mask)
                 y_filt = y_tau[mask]
-                w_filt = (treatment_filt == group).astype(int)
+                w_filt = (treatment_filt_np == group).astype(int)
                 p_filt = cur_p[group][mask]
                 mu_t = self.models_mu_t[group][ifold].predict(X_filt)
                 mu_c = self.models_mu_c[ifold].predict(X_filt)
@@ -202,45 +178,29 @@ def fit(self, X, treatment, y, p=None, seed=None):
                 self.models_tau[group][ifold].fit(X_filt, dr)
 
     def bootstrap(self, X, treatment, y, p=None, size=10000, rng=None, seed=None):
-        """Runs a single bootstrap with optional deterministic cross-fit seed."""
         if rng is not None:
-            idxs = rng.choice(np.arange(0, X.shape[0]), size=size)
+            idxs = rng.choice(np.arange(0, to_numpy(X).shape[0]), size=size)
         else:
-            idxs = np.random.choice(np.arange(0, X.shape[0]), size=size)
-        X_b = X[idxs]
-
-        if p is not None:
-            p_b = {group: _p[idxs] for group, _p in p.items()}
-        else:
-            p_b = None
-
-        treatment_b = treatment[idxs]
-        y_b = y[idxs]
+            idxs = np.random.choice(np.arange(0, to_numpy(X).shape[0]), size=size)
+        X_b = filter_index(X, idxs)
+        p_b = {group: _p[idxs] for group, _p in p.items()} if p is not None else None
+        treatment_b = filter_index(treatment, idxs)
+        y_b = to_numpy(y)[idxs]
         self.fit(X=X_b, treatment=treatment_b, y=y_b, p=p_b, seed=seed)
         return self.predict(X=X, p=p)
 
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        """Predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
-            verbose (bool, optional): whether to output progress logs
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects.
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
-
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         yhat_cs = {}
         yhat_ts = {}
 
         for i, group in enumerate(self.t_groups):
-            models_tau = self.models_tau[group]
-            _te = np.r_[[model.predict(X) for model in models_tau]].mean(axis=0)
+            _te = np.r_[[model.predict(X) for model in self.models_tau[group]]].mean(
+                axis=0
+            )
             te[:, i] = np.ravel(_te)
             yhat_cs[group] = np.r_[
                 [model.predict(X) for model in self.models_mu_c]
@@ -250,10 +210,11 @@ def predict(
             ].mean(axis=0)
 
             if (y is not None) and (treatment is not None) and verbose:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                y_filt = y[mask]
-                w = (treatment_filt == group).astype(int)
+                treatment_np = to_numpy(treatment)
+                mask = (treatment_np == group) | (treatment_np == self.control_name)
+                treatment_filt_np = treatment_np[mask]
+                y_filt = to_numpy(filter_mask(y, mask))
+                w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
                 yhat[w == 0] = yhat_cs[group][mask][w == 0]
@@ -280,27 +241,6 @@ def fit_predict(
         verbose=True,
         seed=None,
     ):
-        """Fit the treatment effect and outcome models of the R learner and predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            return_ci (bool): whether to return confidence intervals
-            n_bootstraps (int): number of bootstrap iterations
-            bootstrap_size (int): number of samples per bootstrap
-            return_components (bool, optional): whether to return outcome for treatment and control seperately
-            verbose (str): whether to output progress logs
-            seed (int): random seed for cross-fitting
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment]
-                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
-                UB [n_samples, n_treatment]
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         self.fit(X, treatment, y, p, seed)
 
         if p is None:
@@ -308,12 +248,9 @@ def fit_predict(
 
         check_p_conditions(p, self.t_groups)
         if isinstance(p, (np.ndarray, pd.Series)):
-            treatment_name = self.t_groups[0]
-            p = {treatment_name: convert_pd_to_np(p)}
+            p = {self.t_groups[0]: to_numpy(p)}
         elif isinstance(p, dict):
-            p = {
-                treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()
-            }
+            p = {k: to_numpy(v) for k, v in p.items()}
 
         te = self.predict(
             X, treatment=treatment, y=y, return_components=return_components
@@ -322,15 +259,18 @@ def fit_predict(
         if not return_ci:
             return te
         else:
+            X_np = to_numpy(X)
+            treatment_np = to_numpy(treatment)
+            y_np = to_numpy(y)
+
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_mu_c_global = deepcopy(self.models_mu_c)
             models_mu_t_global = deepcopy(self.models_mu_t)
             models_tau_global = deepcopy(self.models_tau)
             te_bootstraps = np.zeros(
-                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
             )
-            # seed controls both bootstrap resampling and cross-fit randomness.
             rng = np.random.default_rng(seed) if seed is not None else None
 
             logger.info("Bootstrap Confidence Intervals")
@@ -341,9 +281,9 @@ def fit_predict(
                     else None
                 )
                 te_b = self.bootstrap(
-                    X,
-                    treatment,
-                    y,
+                    X_np,
+                    treatment_np,
+                    y_np,
                     p,
                     size=bootstrap_size,
                     rng=rng,
@@ -356,7 +296,6 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_mu_c = deepcopy(models_mu_c_global)
@@ -377,23 +316,6 @@ def estimate_ate(
         seed=None,
         pretrain=False,
     ):
-        """Estimate the Average Treatment Effect (ATE).
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            bootstrap_ci (bool): whether run bootstrap for confidence intervals
-            n_bootstraps (int): number of bootstrap iterations
-            bootstrap_size (int): number of samples per bootstrap
-            seed (int): random seed for cross-fitting
-            pretrain (bool): whether a model has been fit, default False.
-        Returns:
-            The mean and confidence interval (LB, UB) of the ATE estimate.
-        """
         if pretrain:
             te, yhat_cs, yhat_ts = self.predict(
                 X, treatment, y, p, return_components=True
@@ -402,19 +324,18 @@ def estimate_ate(
             te, yhat_cs, yhat_ts = self.fit_predict(
                 X, treatment, y, p, return_components=True, seed=seed
             )
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
+
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
 
         if p is None:
             p = self.propensity
         else:
             check_p_conditions(p, self.t_groups)
         if isinstance(p, (np.ndarray, pd.Series)):
-            treatment_name = self.t_groups[0]
-            p = {treatment_name: convert_pd_to_np(p)}
+            p = {self.t_groups[0]: to_numpy(p)}
         elif isinstance(p, dict):
-            p = {
-                treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()
-            }
+            p = {k: to_numpy(v) for k, v in p.items()}
 
         ate = np.zeros(self.t_groups.shape[0])
         ate_lb = np.zeros(self.t_groups.shape[0])
@@ -422,18 +343,15 @@ def estimate_ate(
 
         for i, group in enumerate(self.t_groups):
             _ate = te[:, i].mean()
-
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = treatment_np[mask]
             w = (treatment_filt == group).astype(int)
             prob_treatment = float(sum(w)) / w.shape[0]
 
             yhat_c = yhat_cs[group][mask]
             yhat_t = yhat_ts[group][mask]
-            y_filt = y[mask]
+            y_filt = y_np[mask]
 
-            # SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.
-            # "Recent Developments in the Econometrics of Program Evaluation." Journal of Economic Literature
             se = np.sqrt(
                 (
                     (y_filt[w == 0] - yhat_c[w == 0]).var() / (1 - prob_treatment)
@@ -453,6 +371,7 @@ def estimate_ate(
         if not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
+            X_np = to_numpy(X)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_mu_c_global = deepcopy(self.models_mu_c)
@@ -461,7 +380,6 @@ def estimate_ate(
 
             logger.info("Bootstrap Confidence Intervals for ATE")
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
-            # seed controls both bootstrap resampling and cross-fit randomness.
             rng = np.random.default_rng(seed) if seed is not None else None
 
             for n in tqdm(range(n_bootstraps)):
@@ -471,9 +389,9 @@ def estimate_ate(
                     else None
                 )
                 cate_b = self.bootstrap(
-                    X,
-                    treatment,
-                    y,
+                    X_np,
+                    treatment_np,
+                    y_np,
                     p,
                     size=bootstrap_size,
                     rng=rng,
@@ -488,7 +406,6 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_mu_c = deepcopy(models_mu_c_global)
@@ -498,10 +415,6 @@ def estimate_ate(
 
 
 class BaseDRRegressor(BaseDRLearner):
-    """
-    A parent class for DR-learner regressor classes.
-    """
-
     def __init__(
         self,
         learner=None,
@@ -511,17 +424,6 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
-        """Initialize an DR-learner regressor.
-
-        Args:
-            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
-                groups
-            control_outcome_learner (optional): a model to estimate outcomes in the control group
-            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
-            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         super().__init__(
             learner=learner,
             control_outcome_learner=control_outcome_learner,
@@ -533,10 +435,6 @@ def __init__(
 
 
 class BaseDRClassifier(BaseDRLearner):
-    """
-    A parent class for DR-learner classifier classes.
-    """
-
     def __init__(
         self,
         learner=None,
@@ -546,20 +444,6 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
-        """Initialize a DR-learner classifier.
-
-        Args:
-            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
-                groups. Should have a predict_proba() method for outcome models.
-            control_outcome_learner (optional): a model to estimate outcomes in the control group.
-                Should have a predict_proba() method.
-            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group.
-                Should have a predict_proba() method.
-            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group.
-                Should be a regressor.
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         super().__init__(
             learner=learner,
             control_outcome_learner=control_outcome_learner,
@@ -572,35 +456,15 @@ def __init__(
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        """Predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector. Used for computing
-                classification metrics when y is also provided.
-            y (np.array or pd.Series, optional): an outcome vector. Used for computing
-                classification metrics when treatment is also provided.
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1). Currently not used in prediction but kept for API consistency.
-            return_components (bool, optional): whether to return outcome probabilities for treatment and control
-                groups separately. Defaults to False.
-            verbose (bool, optional): whether to output progress logs. Defaults to True.
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects.
-            If return_components is True, also returns:
-                - dict: Predicted probabilities for the control group (yhat_cs).
-                - dict: Predicted probabilities for the treatment group (yhat_ts).
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
-
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         yhat_cs = {}
         yhat_ts = {}
 
         for i, group in enumerate(self.t_groups):
-            models_tau = self.models_tau[group]
-            _te = np.r_[[model.predict(X) for model in models_tau]].mean(axis=0)
+            _te = np.r_[[model.predict(X) for model in self.models_tau[group]]].mean(
+                axis=0
+            )
             te[:, i] = np.ravel(_te)
             yhat_cs[group] = np.r_[
                 [model.predict_proba(X)[:, 1] for model in self.models_mu_c]
@@ -610,10 +474,11 @@ def predict(
             ].mean(axis=0)
 
             if (y is not None) and (treatment is not None) and verbose:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                y_filt = y[mask]
-                w = (treatment_filt == group).astype(int)
+                treatment_np = to_numpy(treatment)
+                mask = (treatment_np == group) | (treatment_np == self.control_name)
+                treatment_filt_np = treatment_np[mask]
+                y_filt = to_numpy(filter_mask(y, mask))
+                w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
                 yhat[w == 0] = yhat_cs[group][mask][w == 0]
@@ -630,7 +495,6 @@ def predict(
 
 class XGBDRRegressor(BaseDRRegressor):
     def __init__(self, ate_alpha=0.05, control_name=0, *args, **kwargs):
-        """Initialize a DR-learner with two XGBoost models."""
         super().__init__(
             learner=XGBRegressor(*args, **kwargs),
             ate_alpha=ate_alpha,
diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py
index 53563cce..30b79d75 100644
--- a/causalml/inference/meta/rlearner.py
+++ b/causalml/inference/meta/rlearner.py
@@ -9,8 +9,9 @@
 from causalml.inference.meta.base import BaseLearner
 from causalml.inference.meta.utils import (
     check_treatment_vector,
+    filter_mask,
+    to_numpy,
     get_xgboost_objective_metric,
-    convert_pd_to_np,
     get_weighted_variance,
 )
 from causalml.propensity import ElasticNetPropensityModel
@@ -19,12 +20,7 @@
 
 
 class BaseRLearner(BaseLearner):
-    """A parent class for R-learner classes.
-
-    An R-learner estimates treatment effects with two machine learning models and the propensity score.
-
-    Details of R-learner are available at `Nie and Wager (2019) <https://arxiv.org/abs/1712.04912>`_.
-    """
+    """A parent class for R-learner classes."""
 
     def __init__(
         self,
@@ -38,22 +34,6 @@ def __init__(
         random_state=None,
         cv_n_jobs=-1,
     ):
-        """Initialize an R-learner.
-
-        Args:
-            learner (optional): a model to estimate outcomes and treatment effects
-            outcome_learner (optional): a model to estimate outcomes
-            effect_learner (optional): a model to estimate treatment effects. It needs to take `sample_weight` as an
-                input argument for `fit()`
-            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
-                be used by default.
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-            n_fold (int, optional): the number of cross validation folds for outcome_learner
-            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
-            cv_n_jobs (int, optional): number of parallel jobs to run for cross_val_predict. -1 means using all
-                processors
-        """
         assert (learner is not None) or (
             (outcome_learner is not None) and (effect_learner is not None)
         )
@@ -69,11 +49,9 @@ def __init__(
 
         self.ate_alpha = ate_alpha
         self.control_name = control_name
-
         self.random_state = random_state
         self.cv = KFold(n_splits=n_fold, shuffle=True, random_state=random_state)
         self.cv_n_jobs = cv_n_jobs
-
         self.propensity = None
         self.propensity_model = None
 
@@ -86,31 +64,21 @@ def __repr__(self):
         )
 
     def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
-        """Fit the treatment effect and outcome models of the R learner.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
-                weight of each observation for `effect_learner`. If None, it assumes equal weight.
-            verbose (bool, optional): whether to output progress logs
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
         if sample_weight is not None:
             assert len(sample_weight) == len(
                 y
             ), "Data length must be equal for sample_weight and the input data"
-            sample_weight = convert_pd_to_np(sample_weight)
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
+            sample_weight = to_numpy(sample_weight)
+
+        self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
 
         if p is None:
-            self._set_propensity_models(X=X, treatment=treatment, y=y)
+            self._set_propensity_models(X=to_numpy(X), treatment=treatment_np, y=y_np)
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -122,27 +90,32 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
 
         if verbose:
             logger.info("generating out-of-fold CV outcome estimates")
-        yhat = cross_val_predict(self.model_mu, X, y, cv=self.cv, n_jobs=self.cv_n_jobs)
+        # sklearn >= 1.6 accepts DataFrames natively — pass X as-is
+        yhat = cross_val_predict(
+            self.model_mu, X, y_np, cv=self.cv, n_jobs=self.cv_n_jobs
+        )
 
         for group in self.t_groups:
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            X_filt = X[mask]
-            y_filt = y[mask]
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = filter_mask(treatment, mask)
+            X_filt = filter_mask(X, mask)
+            y_filt = y_np[mask]
             yhat_filt = yhat[mask]
             p_filt = p[group][mask]
-            w = (treatment_filt == group).astype(int)
+            w = (to_numpy(treatment_filt) == group).astype(int)
 
             weight = (w - p_filt) ** 2
             diff_c = y_filt[w == 0] - yhat_filt[w == 0]
             diff_t = y_filt[w == 1] - yhat_filt[w == 1]
             if sample_weight is not None:
                 sample_weight_filt = sample_weight[mask]
-                sample_weight_filt_c = sample_weight_filt[w == 0]
-                sample_weight_filt_t = sample_weight_filt[w == 1]
-                self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
-                self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
-                weight *= sample_weight_filt  # update weight
+                self.vars_c[group] = get_weighted_variance(
+                    diff_c, sample_weight_filt[w == 0]
+                )
+                self.vars_t[group] = get_weighted_variance(
+                    diff_t, sample_weight_filt[w == 1]
+                )
+                weight *= sample_weight_filt
             else:
                 self.vars_c[group] = diff_c.var()
                 self.vars_t[group] = diff_t.var()
@@ -158,20 +131,10 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             )
 
     def predict(self, X, p=None):
-        """Predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects.
-        """
-        X = convert_pd_to_np(X)
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
-            dhat = self.models_tau[group].predict(X)
-            te[:, i] = dhat
-
+            te[:, i] = self.models_tau[group].predict(X)
         return te
 
     def fit_predict(
@@ -186,39 +149,22 @@ def fit_predict(
         bootstrap_size=10000,
         verbose=True,
     ):
-        """Fit the treatment effect and outcome models of the R learner and predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
-                weight of each observation for `effect_learner`. If None, it assumes equal weight.
-            return_ci (bool): whether to return confidence intervals
-            n_bootstraps (int): number of bootstrap iterations
-            bootstrap_size (int): number of samples per bootstrap
-            verbose (bool): whether to output progress logs
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
-                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
-                UB [n_samples, n_treatment]
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         self.fit(X, treatment, y, p, sample_weight, verbose=verbose)
         te = self.predict(X)
 
         if not return_ci:
             return te
         else:
+            X_np = to_numpy(X)
+            treatment_np = to_numpy(treatment)
+            y_np = to_numpy(y)
+
             t_groups_global = self.t_groups
             _classes_global = self._classes
             model_mu_global = deepcopy(self.model_mu)
             models_tau_global = deepcopy(self.models_tau)
             te_bootstraps = np.zeros(
-                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
             )
 
             logger.info("Bootstrap Confidence Intervals")
@@ -227,7 +173,7 @@ def fit_predict(
                     p = self.propensity
                 else:
                     p = self._format_p(p, self.t_groups)
-                te_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
+                te_b = self.bootstrap(X_np, treatment_np, y_np, p, size=bootstrap_size)
                 te_bootstraps[:, :, i] = te_b
 
             te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
@@ -235,7 +181,6 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.model_mu = deepcopy(model_mu_global)
@@ -255,30 +200,14 @@ def estimate_ate(
         bootstrap_size=10000,
         pretrain=False,
     ):
-        """Estimate the Average Treatment Effect (ATE).
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): only needed when pretrain=False, a treatment vector
-            y (np.array or pd.Series):only needed when pretrain=False, an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
-                weight of each observation for `effect_learner`. If None, it assumes equal weight.
-            bootstrap_ci (bool): whether run bootstrap for confidence intervals
-            n_bootstraps (int): number of bootstrap iterations
-            bootstrap_size (int): number of samples per bootstrap
-            pretrain (bool): whether a model has been fit, default False.
-        Returns:
-            The mean and confidence interval (LB, UB) of the ATE estimate.
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
+        treatment_np = to_numpy(treatment)
+        X_np = to_numpy(X)
+
         if pretrain:
             te = self.predict(X, p)
         else:
-            if not len(treatment) or not len(y):
-                raise ValueError("treatmeng and y must be provided when pretrain=False")
+            if not len(treatment_np) or not len(to_numpy(y)):
+                raise ValueError("treatment and y must be provided when pretrain=False")
             te = self.fit_predict(X, treatment, y, p, sample_weight, return_ci=False)
 
         ate = np.zeros(self.t_groups.shape[0])
@@ -286,8 +215,8 @@ def estimate_ate(
         ate_ub = np.zeros(self.t_groups.shape[0])
 
         for i, group in enumerate(self.t_groups):
-            w = (treatment == group).astype(int)
-            prob_treatment = float(sum(w)) / X.shape[0]
+            w = (treatment_np == group).astype(int)
+            prob_treatment = float(sum(w)) / X_np.shape[0]
             _ate = te[:, i].mean()
 
             se = (
@@ -296,7 +225,7 @@ def estimate_ate(
                     + (self.vars_c[group] / (1 - prob_treatment))
                     + te[:, i].var()
                 )
-                / X.shape[0]
+                / X_np.shape[0]
             )
 
             _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
@@ -309,6 +238,7 @@ def estimate_ate(
         if not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
+            y_np = to_numpy(y)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             model_mu_global = deepcopy(self.model_mu)
@@ -322,7 +252,9 @@ def estimate_ate(
                     p = self.propensity
                 else:
                     p = self._format_p(p, self.t_groups)
-                cate_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
+                cate_b = self.bootstrap(
+                    X_np, treatment_np, y_np, p, size=bootstrap_size
+                )
                 ate_bootstraps[:, n] = cate_b.mean(axis=0)
 
             ate_lower = np.percentile(
@@ -332,7 +264,6 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.model_mu = deepcopy(model_mu_global)
@@ -341,10 +272,6 @@ def estimate_ate(
 
 
 class BaseRRegressor(BaseRLearner):
-    """
-    A parent class for R-learner regressor classes.
-    """
-
     def __init__(
         self,
         learner=None,
@@ -356,20 +283,6 @@ def __init__(
         n_fold=5,
         random_state=None,
     ):
-        """Initialize an R-learner regressor.
-
-        Args:
-            learner (optional): a model to estimate outcomes and treatment effects
-            outcome_learner (optional): a model to estimate outcomes
-            effect_learner (optional): a model to estimate treatment effects. It needs to take `sample_weight` as an
-                input argument for `fit()`
-            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
-                be used by default.
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-            n_fold (int, optional): the number of cross validation folds for outcome_learner
-            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
-        """
         super().__init__(
             learner=learner,
             outcome_learner=outcome_learner,
@@ -383,10 +296,6 @@ def __init__(
 
 
 class BaseRClassifier(BaseRLearner):
-    """
-    A parent class for R-learner classifier classes.
-    """
-
     def __init__(
         self,
         outcome_learner=None,
@@ -397,19 +306,6 @@ def __init__(
         n_fold=5,
         random_state=None,
     ):
-        """Initialize an R-learner classifier.
-
-        Args:
-            outcome_learner: a model to estimate outcomes. Should be a classifier.
-            effect_learner: a model to estimate treatment effects. It needs to take `sample_weight` as an
-                input argument for `fit()`. Should be a regressor.
-            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
-                be used by default.
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-            n_fold (int, optional): the number of cross validation folds for outcome_learner
-            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
-        """
         super().__init__(
             learner=None,
             outcome_learner=outcome_learner,
@@ -420,38 +316,27 @@ def __init__(
             n_fold=n_fold,
             random_state=random_state,
         )
-
         if (outcome_learner is None) and (effect_learner is None):
             raise ValueError(
                 "Either the outcome learner or the effect learner must be specified."
             )
 
     def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
-        """Fit the treatment effect and outcome models of the R learner.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
-                weight of each observation for `effect_learner`. If None, it assumes equal weight.
-            verbose (bool, optional): whether to output progress logs
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
         if sample_weight is not None:
             assert len(sample_weight) == len(
                 y
             ), "Data length must be equal for sample_weight and the input data"
-            sample_weight = convert_pd_to_np(sample_weight)
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
+            sample_weight = to_numpy(sample_weight)
+
+        self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
 
         if p is None:
-            self._set_propensity_models(X=X, treatment=treatment, y=y)
+            self._set_propensity_models(X=to_numpy(X), treatment=treatment_np, y=y_np)
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -464,28 +349,30 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
         if verbose:
             logger.info("generating out-of-fold CV outcome estimates")
         yhat = cross_val_predict(
-            self.model_mu, X, y, cv=self.cv, method="predict_proba", n_jobs=-1
+            self.model_mu, X, y_np, cv=self.cv, method="predict_proba", n_jobs=-1
         )[:, 1]
 
         for group in self.t_groups:
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            X_filt = X[mask]
-            y_filt = y[mask]
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = filter_mask(treatment, mask)
+            X_filt = filter_mask(X, mask)
+            y_filt = y_np[mask]
             yhat_filt = yhat[mask]
             p_filt = p[group][mask]
-            w = (treatment_filt == group).astype(int)
+            w = (to_numpy(treatment_filt) == group).astype(int)
 
             weight = (w - p_filt) ** 2
             diff_c = y_filt[w == 0] - yhat_filt[w == 0]
             diff_t = y_filt[w == 1] - yhat_filt[w == 1]
             if sample_weight is not None:
                 sample_weight_filt = sample_weight[mask]
-                sample_weight_filt_c = sample_weight_filt[w == 0]
-                sample_weight_filt_t = sample_weight_filt[w == 1]
-                self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
-                self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
-                weight *= sample_weight_filt  # update weight
+                self.vars_c[group] = get_weighted_variance(
+                    diff_c, sample_weight_filt[w == 0]
+                )
+                self.vars_t[group] = get_weighted_variance(
+                    diff_t, sample_weight_filt[w == 1]
+                )
+                weight *= sample_weight_filt
             else:
                 self.vars_c[group] = diff_c.var()
                 self.vars_t[group] = diff_t.var()
@@ -501,20 +388,10 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             )
 
     def predict(self, X, p=None):
-        """Predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects.
-        """
-        X = convert_pd_to_np(X)
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
-            dhat = self.models_tau[group].predict(X)
-            te[:, i] = dhat
-
+            te[:, i] = self.models_tau[group].predict(X)
         return te
 
 
@@ -530,21 +407,7 @@ def __init__(
         *args,
         **kwargs,
     ):
-        """Initialize an R-learner regressor with XGBoost model using pairwise ranking objective.
-
-        Args:
-            early_stopping: whether or not to use early stopping when fitting effect learner
-            test_size (float, optional): the proportion of the dataset to use as validation set when early stopping is
-                                         enabled
-            early_stopping_rounds (int, optional): validation metric needs to improve at least once in every
-                                                   early_stopping_rounds round(s) to continue training
-            effect_learner_objective (str, optional): the learning objective for the effect learner
-                                                      (default = 'reg:squarederror')
-            effect_learner_n_estimators (int, optional): number of trees to fit for the effect learner (default = 500)
-        """
-
         assert isinstance(random_state, int), "random_state should be int."
-
         objective, metric = get_xgboost_objective_metric(effect_learner_objective)
         self.effect_learner_objective = objective
         self.effect_learner_eval_metric = metric
@@ -553,7 +416,6 @@ def __init__(
         if self.early_stopping:
             self.test_size = test_size
             self.early_stopping_rounds = early_stopping_rounds
-
             effect_learner = XGBRegressor(
                 objective=self.effect_learner_objective,
                 n_estimators=self.effect_learner_n_estimators,
@@ -572,41 +434,28 @@ def __init__(
                 *args,
                 **kwargs,
             )
-
         super().__init__(
             outcome_learner=XGBRegressor(random_state=random_state, *args, **kwargs),
             effect_learner=effect_learner,
         )
 
     def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
-        """Fit the treatment effect and outcome models of the R learner.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            sample_weight (np.array or pd.Series, optional): an array of sample weights indicating the
-                weight of each observation for `effect_learner`. If None, it assumes equal weight.
-            verbose (bool, optional): whether to output progress logs
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
-        # initialize equal sample weight if it's not provided, for simplicity purpose
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
         sample_weight = (
-            convert_pd_to_np(sample_weight)
-            if sample_weight is not None
-            else convert_pd_to_np(np.ones(len(y)))
+            to_numpy(sample_weight) if sample_weight is not None else np.ones(len(y_np))
         )
         assert len(sample_weight) == len(
-            y
+            y_np
         ), "Data length must be equal for sample_weight and the input data"
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
+
+        self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
 
         if p is None:
-            self._set_propensity_models(X=X, treatment=treatment, y=y)
+            self._set_propensity_models(X=to_numpy(X), treatment=treatment_np, y=y_np)
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -618,18 +467,18 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
 
         if verbose:
             logger.info("generating out-of-fold CV outcome estimates")
-        yhat = cross_val_predict(self.model_mu, X, y, cv=self.cv, n_jobs=-1)
+        yhat = cross_val_predict(self.model_mu, X, y_np, cv=self.cv, n_jobs=-1)
 
         for group in self.t_groups:
-            treatment_mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[treatment_mask]
-            w = (treatment_filt == group).astype(int)
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt_np = treatment_np[mask]
+            w = (treatment_filt_np == group).astype(int)
 
-            X_filt = X[treatment_mask]
-            y_filt = y[treatment_mask]
-            yhat_filt = yhat[treatment_mask]
-            p_filt = p[group][treatment_mask]
-            sample_weight_filt = sample_weight[treatment_mask]
+            X_filt = filter_mask(X, mask)
+            y_filt = y_np[mask]
+            yhat_filt = yhat[mask]
+            p_filt = p[group][mask]
+            sample_weight_filt = sample_weight[mask]
 
             if verbose:
                 logger.info(
@@ -662,7 +511,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
                     test_size=self.test_size,
                     random_state=self.random_state,
                 )
-
                 self.models_tau[group].fit(
                     X=X_train_filt,
                     y=(y_train_filt - yhat_train_filt) / (w_train - p_train_filt),
@@ -679,7 +527,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
                     ],
                     verbose=verbose,
                 )
-
             else:
                 self.models_tau[group].fit(
                     X_filt,
@@ -689,7 +536,9 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
 
             diff_c = y_filt[w == 0] - yhat_filt[w == 0]
             diff_t = y_filt[w == 1] - yhat_filt[w == 1]
-            sample_weight_filt_c = sample_weight_filt[w == 0]
-            sample_weight_filt_t = sample_weight_filt[w == 1]
-            self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
-            self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
+            self.vars_c[group] = get_weighted_variance(
+                diff_c, sample_weight_filt[w == 0]
+            )
+            self.vars_t[group] = get_weighted_variance(
+                diff_t, sample_weight_filt[w == 1]
+            )
diff --git a/causalml/inference/meta/slearner.py b/causalml/inference/meta/slearner.py
index 6891f02d..445d2f4d 100644
--- a/causalml/inference/meta/slearner.py
+++ b/causalml/inference/meta/slearner.py
@@ -7,7 +7,13 @@
 from copy import deepcopy
 
 from causalml.inference.meta.base import BaseLearner
-from causalml.inference.meta.utils import check_treatment_vector, convert_pd_to_np
+from causalml.inference.meta.utils import (
+    check_treatment_vector,
+    filter_mask,
+    prepend_column,
+    concat_treatment_col,
+    to_numpy,
+)
 from causalml.metrics import regression_metrics, classification_metrics
 
 logger = logging.getLogger("causalml")
@@ -17,44 +23,24 @@ class StatsmodelsOLS:
     """A sklearn style wrapper class for statsmodels' OLS."""
 
     def __init__(self, cov_type="HC1", alpha=0.05):
-        """Initialize a statsmodels' OLS wrapper class object.
-        Args:
-            cov_type (str, optional): covariance estimator type.
-            alpha (float, optional): the confidence level alpha.
-        """
         self.cov_type = cov_type
         self.alpha = alpha
 
     def fit(self, X, y):
-        """Fit OLS.
-        Args:
-            X (np.matrix): a feature matrix
-            y (np.array): a label vector
-        """
-        # Append ones. The first column is for the treatment indicator.
         X = sm.add_constant(X, prepend=False, has_constant="add")
         self.model = sm.OLS(y, X).fit(cov_type=self.cov_type)
         self.coefficients = self.model.params
         self.conf_ints = self.model.conf_int(alpha=self.alpha)
 
     def predict(self, X):
-        # Append ones. The first column is for the treatment indicator.
         X = sm.add_constant(X, prepend=False, has_constant="add")
         return self.model.predict(X)
 
 
 class BaseSLearner(BaseLearner):
-    """A parent class for S-learner classes.
-    An S-learner estimates treatment effects with one machine learning model.
-    Details of S-learner are available at `Kunzel et al. (2018) <https://arxiv.org/abs/1706.03461>`_.
-    """
+    """A parent class for S-learner classes."""
 
     def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
-        """Initialize an S-learner.
-        Args:
-            learner (optional): a model to estimate the treatment effect
-            control_name (str or int, optional): name of control group
-        """
         if learner is not None:
             self.model = learner
         else:
@@ -66,27 +52,27 @@ def __repr__(self):
         return "{}(model={})".format(self.__class__.__name__, self.model.__repr__())
 
     def fit(self, X, treatment, y, p=None):
-        """Fit the inference model
+        """Fit the inference model.
         Args:
-            X (np.matrix, np.array, or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
+            X (np.matrix, np.array, pd.Dataframe, or pl.DataFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
         """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        treatment_np = to_numpy(treatment)
+        self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
         self.models = {group: deepcopy(self.model) for group in self.t_groups}
 
         for group in self.t_groups:
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            X_filt = X[mask]
-            y_filt = y[mask]
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = filter_mask(treatment, mask)
+            X_filt = filter_mask(X, mask)
+            y_filt = filter_mask(y, mask)
 
-            w = (treatment_filt == group).astype(int)
-            X_new = np.hstack((w.reshape((-1, 1)), X_filt))
+            w = (to_numpy(treatment_filt) == group).astype(int)
+            X_new = concat_treatment_col(w, X_filt)
             self.models[group].fit(X_new, y_filt)
 
     def predict(
@@ -94,35 +80,32 @@ def predict(
     ):
         """Predict treatment effects.
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
+            X (np.matrix, np.array, pd.Dataframe, or pl.DataFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
+            y (np.array, pd.Series, or pl.Series, optional): an outcome vector
             return_components (bool, optional): whether to return outcome for treatment and control seperately
             verbose (bool, optional): whether to output progress logs
         Returns:
             (numpy.ndarray): Predictions of treatment effects.
         """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         yhat_cs = {}
         yhat_ts = {}
 
         for group in self.t_groups:
             model = self.models[group]
 
-            # Build separate arrays for control and treatment to avoid in-place
-            # mutation, which fails when learners like CatBoost set the
-            # writeable flag to False on arrays passed to predict().
-            X_new_c = np.hstack((np.zeros((X.shape[0], 1)), X))
+            X_new_c = prepend_column(0.0, X)
             yhat_cs[group] = model.predict(X_new_c)
 
-            X_new_t = np.hstack((np.ones((X.shape[0], 1)), X))
+            X_new_t = prepend_column(1.0, X)
             yhat_ts[group] = model.predict(X_new_t)
 
             if (y is not None) and (treatment is not None) and verbose:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                w = (treatment_filt == group).astype(int)
-                y_filt = y[mask]
+                treatment_np = to_numpy(treatment)
+                mask = (treatment_np == group) | (treatment_np == self.control_name)
+                treatment_filt_np = treatment_np[mask]
+                w = (treatment_filt_np == group).astype(int)
+                y_filt = to_numpy(filter_mask(y, mask))
 
                 yhat = np.zeros_like(y_filt, dtype=float)
                 yhat[w == 0] = yhat_cs[group][mask][w == 0]
@@ -131,7 +114,8 @@ def predict(
                 logger.info("Error metrics for group {}".format(group))
                 regression_metrics(y_filt, yhat, w)
 
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
@@ -152,37 +136,26 @@ def fit_predict(
         return_components=False,
         verbose=True,
     ):
-        """Fit the inference model of the S learner and predict treatment effects.
-        Args:
-            X (np.matrix, np.array, or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            return_ci (bool, optional): whether to return confidence intervals
-            n_bootstraps (int, optional): number of bootstrap iterations
-            bootstrap_size (int, optional): number of samples per bootstrap
-            return_components (bool, optional): whether to return outcome for treatment and control seperately
-            verbose (bool, optional): whether to output progress logs
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
-                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
-                UB [n_samples, n_treatment]
-        """
         self.fit(X, treatment, y)
         te = self.predict(X, treatment, y, return_components=return_components)
 
         if not return_ci:
             return te
         else:
+            X_np = to_numpy(X)
+            treatment_np = to_numpy(treatment)
+            y_np = to_numpy(y)
+
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_global = deepcopy(self.models)
             te_bootstraps = np.zeros(
-                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
             )
 
             logger.info("Bootstrap Confidence Intervals")
             for i in tqdm(range(n_bootstraps)):
-                te_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
+                te_b = self.bootstrap(X_np, treatment_np, y_np, size=bootstrap_size)
                 te_bootstraps[:, :, i] = te_b
 
             te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
@@ -190,7 +163,6 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models = deepcopy(models_global)
@@ -209,22 +181,6 @@ def estimate_ate(
         bootstrap_size=10000,
         pretrain=False,
     ):
-        """Estimate the Average Treatment Effect (ATE).
-
-        Args:
-            X (np.matrix, np.array, or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            return_ci (bool, optional): whether to return confidence intervals
-            bootstrap_ci (bool): whether to return confidence intervals
-            n_bootstraps (int): number of bootstrap iterations
-            bootstrap_size (int): number of samples per bootstrap
-            pretrain (bool): whether a model has been fit, default False.
-        Returns:
-            The mean and confidence interval (LB, UB) of the ATE estimate.
-        """
-
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         if pretrain:
             te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True)
         else:
@@ -232,6 +188,9 @@ def estimate_ate(
                 X, treatment, y, return_components=True
             )
 
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
         ate = np.zeros(self.t_groups.shape[0])
         ate_lb = np.zeros(self.t_groups.shape[0])
         ate_ub = np.zeros(self.t_groups.shape[0])
@@ -239,9 +198,9 @@ def estimate_ate(
         for i, group in enumerate(self.t_groups):
             _ate = te[:, i].mean()
 
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            y_filt = y[mask]
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = treatment_np[mask]
+            y_filt = y_np[mask]
             w = (treatment_filt == group).astype(int)
             prob_treatment = float(sum(w)) / w.shape[0]
 
@@ -269,6 +228,7 @@ def estimate_ate(
         elif return_ci and not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
+            X_np = to_numpy(X)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_global = deepcopy(self.models)
@@ -277,7 +237,7 @@ def estimate_ate(
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
 
             for n in tqdm(range(n_bootstraps)):
-                ate_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
+                ate_b = self.bootstrap(X_np, treatment_np, y_np, size=bootstrap_size)
                 ate_bootstraps[:, n] = ate_b.mean(axis=0)
 
             ate_lower = np.percentile(
@@ -287,7 +247,6 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models = deepcopy(models_global)
@@ -296,33 +255,14 @@ def estimate_ate(
 
 
 class BaseSRegressor(BaseSLearner):
-    """
-    A parent class for S-learner regressor classes.
-    """
-
     def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
-        """Initialize an S-learner regressor.
-        Args:
-            learner (optional): a model to estimate the treatment effect
-            control_name (str or int, optional): name of control group
-        """
         super().__init__(
             learner=learner, ate_alpha=ate_alpha, control_name=control_name
         )
 
 
 class BaseSClassifier(BaseSLearner):
-    """
-    A parent class for S-learner classifier classes.
-    """
-
     def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
-        """Initialize an S-learner classifier.
-        Args:
-            learner (optional): a model to estimate the treatment effect.
-                Should have a predict_proba() method.
-            control_name (str or int, optional): name of control group
-        """
         super().__init__(
             learner=learner, ate_alpha=ate_alpha, control_name=control_name
         )
@@ -330,37 +270,24 @@ def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        """Predict treatment effects.
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
-            return_components (bool, optional): whether to return outcome for treatment and control seperately
-            verbose (bool, optional): whether to output progress logs
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects.
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         yhat_cs = {}
         yhat_ts = {}
 
         for group in self.t_groups:
             model = self.models[group]
 
-            # Build separate arrays for control and treatment to avoid in-place
-            # mutation, which fails when learners like CatBoost set the
-            # writeable flag to False on arrays passed to predict().
-            X_new_c = np.hstack((np.zeros((X.shape[0], 1)), X))
+            X_new_c = prepend_column(0.0, X)
             yhat_cs[group] = model.predict_proba(X_new_c)[:, 1]
 
-            X_new_t = np.hstack((np.ones((X.shape[0], 1)), X))
+            X_new_t = prepend_column(1.0, X)
             yhat_ts[group] = model.predict_proba(X_new_t)[:, 1]
 
             if y is not None and (treatment is not None) and verbose:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                w = (treatment_filt == group).astype(int)
-                y_filt = y[mask]
+                treatment_np = to_numpy(treatment)
+                mask = (treatment_np == group) | (treatment_np == self.control_name)
+                treatment_filt_np = treatment_np[mask]
+                w = (treatment_filt_np == group).astype(int)
+                y_filt = to_numpy(filter_mask(y, mask))
 
                 yhat = np.zeros_like(y_filt, dtype=float)
                 yhat[w == 0] = yhat_cs[group][mask][w == 0]
@@ -369,7 +296,8 @@ def predict(
                 logger.info("Error metrics for group {}".format(group))
                 classification_metrics(y_filt, yhat, w)
 
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
@@ -381,23 +309,9 @@ def predict(
 
 class LRSRegressor(BaseSRegressor):
     def __init__(self, ate_alpha=0.05, control_name=0):
-        """Initialize an S-learner with a linear regression model.
-        Args:
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         super().__init__(StatsmodelsOLS(alpha=ate_alpha), ate_alpha, control_name)
 
     def estimate_ate(self, X, treatment, y, p=None, pretrain=False):
-        """Estimate the Average Treatment Effect (ATE).
-        Args:
-            X (np.matrix, np.array, or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-        Returns:
-            The mean and confidence interval (LB, UB) of the ATE estimate.
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         if not pretrain:
             self.fit(X, treatment, y)
 
diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py
index 04ca796f..b80386d7 100644
--- a/causalml/inference/meta/tlearner.py
+++ b/causalml/inference/meta/tlearner.py
@@ -15,7 +15,12 @@
 from xgboost import XGBRegressor
 
 from causalml.inference.meta.base import BaseLearner
-from causalml.inference.meta.utils import check_treatment_vector, convert_pd_to_np
+from causalml.inference.meta.utils import (
+    _POLARS_AVAILABLE,
+    check_treatment_vector,
+    filter_mask,
+    to_numpy,
+)
 from causalml.metrics import regression_metrics, classification_metrics
 
 logger = logging.getLogger("causalml")
@@ -73,27 +78,33 @@ def fit(self, X, treatment, y, p=None):
         """Fit the inference model
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
+            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
+            treatment (np.array or pd.Series or pl.Series): a treatment vector
+            y (np.array or pd.Series or pl.Series): an outcome vector
         """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        self.t_groups = np.unique(
+            to_numpy(treatment)[to_numpy(treatment) != self.control_name]
+        )
         self.t_groups.sort()
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
         self.models_c = {group: deepcopy(self.model_c) for group in self.t_groups}
         self.models_t = {group: deepcopy(self.model_t) for group in self.t_groups}
 
+        treatment_np = to_numpy(treatment)
         for group in self.t_groups:
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            X_filt = X[mask]
-            y_filt = y[mask]
-            w = (treatment_filt == group).astype(int)
-
-            self.models_c[group].fit(X_filt[w == 0], y_filt[w == 0])
-            self.models_t[group].fit(X_filt[w == 1], y_filt[w == 1])
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = filter_mask(treatment, mask)
+            X_filt = filter_mask(X, mask)
+            y_filt = filter_mask(y, mask)
+            w = (to_numpy(treatment_filt) == group).astype(int)
+
+            self.models_c[group].fit(
+                filter_mask(X_filt, w == 0), filter_mask(y_filt, w == 0)
+            )
+            self.models_t[group].fit(
+                filter_mask(X_filt, w == 1), filter_mask(y_filt, w == 1)
+            )
 
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
@@ -101,15 +112,21 @@ def predict(
         """Predict treatment effects.
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
+            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
+            treatment (np.array or pd.Series or pl.Series, optional): a treatment vector
+            y (np.array or pd.Series or pl.Series, optional): an outcome vector
             return_components (bool, optional): whether to return outcome for treatment and control seperately
             verbose (bool, optional): whether to output progress logs
         Returns:
             (numpy.ndarray): Predictions of treatment effects.
         """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
+
+        if _POLARS_AVAILABLE:
+            import polars as pl
+
+            if isinstance(X, pl.LazyFrame):
+                X = X.collect()
+
         yhat_cs = {}
         yhat_ts = {}
 
@@ -120,10 +137,11 @@ def predict(
             yhat_ts[group] = model_t.predict(X)
 
             if (y is not None) and (treatment is not None) and verbose:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                y_filt = y[mask]
-                w = (treatment_filt == group).astype(int)
+                treatment_np = to_numpy(treatment)
+                mask = (treatment_np == group) | (treatment_np == self.control_name)
+                treatment_filt_np = treatment_np[mask]
+                y_filt = to_numpy(filter_mask(y, mask))
+                w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
                 yhat[w == 0] = yhat_cs[group][mask][w == 0]
@@ -132,7 +150,8 @@ def predict(
                 logger.info("Error metrics for group {}".format(group))
                 regression_metrics(y_filt, yhat, w)
 
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
@@ -156,9 +175,9 @@ def fit_predict(
         """Fit the inference model of the T learner and predict treatment effects.
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
+            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
+            treatment (np.array or pd.Series or pl.Series): a treatment vector
+            y (np.array or pd.Series or pl.Series): an outcome vector
             return_ci (bool): whether to return confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -169,24 +188,27 @@ def fit_predict(
                 If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
                 UB [n_samples, n_treatment]
         """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         self.fit(X, treatment, y)
         te = self.predict(X, treatment, y, return_components=return_components)
 
         if not return_ci:
             return te
         else:
+            X_np = to_numpy(X)
+            treatment_np = to_numpy(treatment)
+            y_np = to_numpy(y)
+
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_c_global = deepcopy(self.models_c)
             models_t_global = deepcopy(self.models_t)
             te_bootstraps = np.zeros(
-                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
             )
 
             logger.info("Bootstrap Confidence Intervals")
             for i in tqdm(range(n_bootstraps)):
-                te_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
+                te_b = self.bootstrap(X_np, treatment_np, y_np, size=bootstrap_size)
                 te_bootstraps[:, :, i] = te_b
 
             te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
@@ -194,7 +216,6 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_c = deepcopy(models_c_global)
@@ -216,17 +237,16 @@ def estimate_ate(
         """Estimate the Average Treatment Effect (ATE).
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
+            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
+            treatment (np.array or pd.Series or pl.Series): a treatment vector
+            y (np.array or pd.Series or pl.Series): an outcome vector
             bootstrap_ci (bool): whether to return confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
+            pretrain (bool): whether a model has been fit, default False.
         Returns:
             The mean and confidence interval (LB, UB) of the ATE estimate.
-            pretrain (bool): whether a model has been fit, default False.
         """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         if pretrain:
             te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True)
         else:
@@ -234,6 +254,10 @@ def estimate_ate(
                 X, treatment, y, return_components=True
             )
 
+        # work in numpy for the ATE calculation
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
         ate = np.zeros(self.t_groups.shape[0])
         ate_lb = np.zeros(self.t_groups.shape[0])
         ate_ub = np.zeros(self.t_groups.shape[0])
@@ -241,9 +265,9 @@ def estimate_ate(
         for i, group in enumerate(self.t_groups):
             _ate = te[:, i].mean()
 
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            y_filt = y[mask]
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = treatment_np[mask]
+            y_filt = y_np[mask]
             w = (treatment_filt == group).astype(int)
             prob_treatment = float(sum(w)) / w.shape[0]
 
@@ -269,6 +293,7 @@ def estimate_ate(
         if not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
+            X_np = to_numpy(X)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_c_global = deepcopy(self.models_c)
@@ -278,7 +303,7 @@ def estimate_ate(
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
 
             for n in tqdm(range(n_bootstraps)):
-                ate_b = self.bootstrap(X, treatment, y, size=bootstrap_size)
+                ate_b = self.bootstrap(X_np, treatment_np, y_np, size=bootstrap_size)
                 ate_bootstraps[:, n] = ate_b.mean(axis=0)
 
             ate_lower = np.percentile(
@@ -288,7 +313,6 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_c = deepcopy(models_c_global)
@@ -298,9 +322,7 @@ def estimate_ate(
 
 
 class BaseTRegressor(BaseTLearner):
-    """
-    A parent class for T-learner regressor classes.
-    """
+    """A parent class for T-learner regressor classes."""
 
     def __init__(
         self,
@@ -310,15 +332,6 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
-        """Initialize a T-learner regressor.
-
-        Args:
-            learner (model): a model to estimate control and treatment outcomes.
-            control_learner (model, optional): a model to estimate control outcomes
-            treatment_learner (model, optional): a model to estimate treatment outcomes
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         super().__init__(
             learner=learner,
             control_learner=control_learner,
@@ -329,9 +342,7 @@ def __init__(
 
 
 class BaseTClassifier(BaseTLearner):
-    """
-    A parent class for T-learner classifier classes.
-    """
+    """A parent class for T-learner classifier classes."""
 
     def __init__(
         self,
@@ -341,15 +352,6 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
-        """Initialize a T-learner classifier.
-
-        Args:
-            learner (model): a model to estimate control and treatment outcomes.
-            control_learner (model, optional): a model to estimate control outcomes
-            treatment_learner (model, optional): a model to estimate treatment outcomes
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         super().__init__(
             learner=learner,
             control_learner=control_learner,
@@ -364,13 +366,19 @@ def predict(
         """Predict treatment effects.
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
+            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
+            treatment (np.array or pd.Series or pl.Series, optional): a treatment vector
+            y (np.array or pd.Series or pl.Series, optional): an outcome vector
             verbose (bool, optional): whether to output progress logs
         Returns:
             (numpy.ndarray): Predictions of treatment effects.
         """
+        if _POLARS_AVAILABLE:
+            import polars as pl
+
+            if isinstance(X, pl.LazyFrame):
+                X = X.collect()
+
         yhat_cs = {}
         yhat_ts = {}
 
@@ -381,10 +389,11 @@ def predict(
             yhat_ts[group] = model_t.predict_proba(X)[:, 1]
 
             if (y is not None) and (treatment is not None) and verbose:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                y_filt = y[mask]
-                w = (treatment_filt == group).astype(int)
+                treatment_np = to_numpy(treatment)
+                mask = (treatment_np == group) | (treatment_np == self.control_name)
+                treatment_filt_np = treatment_np[mask]
+                y_filt = to_numpy(filter_mask(y, mask))
+                w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
                 yhat[w == 0] = yhat_cs[group][mask][w == 0]
@@ -393,7 +402,8 @@ def predict(
                 logger.info("Error metrics for group {}".format(group))
                 classification_metrics(y_filt, yhat, w)
 
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
diff --git a/causalml/inference/meta/utils.py b/causalml/inference/meta/utils.py
index be0ceca7..9aeddb85 100644
--- a/causalml/inference/meta/utils.py
+++ b/causalml/inference/meta/utils.py
@@ -4,7 +4,9 @@
 from packaging import version
 from xgboost import __version__ as xgboost_version
 
+# ---------------------------------------------------------------------------
 # Optional Polars import
+# ---------------------------------------------------------------------------
 try:
     import polars as pl
 
@@ -14,88 +16,179 @@
     _POLARS_AVAILABLE = False
 
 
-def _is_polars_dataframe(obj):
-    """Return True if *obj* is a polars DataFrame or LazyFrame."""
-    if not _POLARS_AVAILABLE:
-        return False
-    return isinstance(obj, (pl.DataFrame, pl.LazyFrame))
+# ---------------------------------------------------------------------------
+# Native DataFrame helpers
+# ---------------------------------------------------------------------------
 
 
-def _is_polars_series(obj):
-    """Return True if *obj* is a polars Series."""
-    if not _POLARS_AVAILABLE:
-        return False
-    return isinstance(obj, pl.Series)
+def filter_mask(obj, mask):
+    """Filter rows by a boolean mask.
 
+    Works natively for numpy arrays, pandas DataFrames/Series, and
+    polars DataFrames/Series without materialising unnecessary copies.
 
-def _polars_to_numpy(obj):
-    """Convert a polars DataFrame, LazyFrame, or Series to a NumPy array.
+    Args:
+        obj: numpy array, pd.DataFrame, pd.Series, pl.DataFrame, or pl.Series.
+        mask: boolean array or Series of the same length as obj.
 
-    - ``pl.LazyFrame`` is collected first (implicit ``.collect()``).
-    - A single-column ``pl.DataFrame`` is squeezed to a 1-D array to match
-      the behaviour of ``pd.Series.to_numpy()``.
-    - A multi-column ``pl.DataFrame`` is returned as a 2-D array.
+    Returns:
+        Filtered object of the same type as input.
     """
-    if isinstance(obj, pl.LazyFrame):
-        obj = obj.collect()
+    if obj is None:
+        return None
+    if isinstance(obj, (pd.DataFrame, pd.Series)):
+        if isinstance(mask, pd.Series):
+            return obj.loc[mask]
+        # numpy boolean array — reset index to avoid alignment issues
+        return obj.loc[np.asarray(mask, dtype=bool)]
+    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.LazyFrame, pl.Series)):
+        if isinstance(obj, pl.LazyFrame):
+            obj = obj.collect()
+        if isinstance(mask, pl.Series):
+            return obj.filter(mask)
+        return obj.filter(pl.Series(np.asarray(mask, dtype=bool)))
+    # numpy / anything else
+    return obj[np.asarray(mask, dtype=bool)]
+
+
+def filter_index(obj, indices):
+    """Filter rows by integer positional indices.
+
+    Used for KFold partition slicing (e.g. in DR-learner).
+
+    Args:
+        obj: numpy array, pd.DataFrame, pd.Series, pl.DataFrame, or pl.Series.
+        indices: integer array of row positions.
+
+    Returns:
+        Filtered object of the same type as input.
+    """
+    if obj is None:
+        return None
+    if isinstance(obj, (pd.DataFrame, pd.Series)):
+        return obj.iloc[indices]
+    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.LazyFrame, pl.Series)):
+        if isinstance(obj, pl.LazyFrame):
+            obj = obj.collect()
+        return obj[indices]
+    return obj[indices]  # numpy
+
+
+def prepend_column(value, X):
+    """Prepend a constant-value column to a feature matrix X.
+
+    Used by the S-learner to prepend the treatment indicator before fitting.
+
+    Args:
+        value (float): scalar value to fill the new column with (0.0 or 1.0).
+        X: numpy array, pd.DataFrame, or pl.DataFrame.
+
+    Returns:
+        Feature matrix with the new column prepended, same type as X.
+    """
+    n = len(X)
+    if isinstance(X, pd.DataFrame):
+        col = pd.DataFrame({"_w": np.full(n, value)}, index=X.index)
+        return pd.concat([col, X], axis=1)
+    if _POLARS_AVAILABLE and isinstance(X, pl.DataFrame):
+        col = pl.Series("_w", np.full(n, value))
+        return X.with_columns(col).select(["_w"] + X.columns)
+    # numpy
+    return np.hstack((np.full((n, 1), value), X))
+
+
+def concat_treatment_col(w, X):
+    """Prepend an array-valued treatment column *w* to feature matrix X.
 
-    if isinstance(obj, pl.DataFrame):
+    Unlike :func:`prepend_column`, this takes an existing 1-D array rather
+    than a scalar. Used by S-learner ``fit()`` to build the augmented matrix.
+
+    Args:
+        w: 1-D numpy array of treatment indicators (0/1).
+        X: numpy array, pd.DataFrame, or pl.DataFrame.
+
+    Returns:
+        Augmented feature matrix with w prepended, same type as X.
+    """
+    if isinstance(X, pd.DataFrame):
+        col = pd.DataFrame({"_w": w}, index=X.index)
+        return pd.concat(
+            [col, X.reset_index(drop=True) if not X.index.equals(col.index) else X],
+            axis=1,
+        )
+    if _POLARS_AVAILABLE and isinstance(X, pl.DataFrame):
+        col = pl.Series("_w", np.asarray(w))
+        return X.with_columns(col).select(["_w"] + X.columns)
+    # numpy
+    return np.hstack((np.asarray(w).reshape(-1, 1), X))
+
+
+def to_numpy(obj):
+    """Convert a pandas or polars object to a NumPy array.
+
+    Use this *only* at boundaries where a third-party library strictly
+    requires numpy (rare — sklearn >= 1.6 and XGBoost >= 3.1 both accept
+    DataFrames natively).
+
+    Args:
+        obj: numpy array, pd.DataFrame, pd.Series, pl.DataFrame,
+             pl.LazyFrame, or pl.Series.  ``None`` is passed through.
+
+    Returns:
+        numpy.ndarray or None.
+    """
+    if obj is None:
+        return None
+    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.LazyFrame, pl.Series)):
+        if isinstance(obj, pl.LazyFrame):
+            obj = obj.collect()
         arr = obj.to_numpy()
-        # Squeeze single-column frames so downstream code gets a 1-D vector
-        if arr.shape[1] == 1:
+        if isinstance(obj, pl.DataFrame) and arr.ndim == 2 and arr.shape[1] == 1:
             return arr.ravel()
         return arr
-
-    if isinstance(obj, pl.Series):
+    if hasattr(obj, "to_numpy"):
         return obj.to_numpy()
-
-    raise TypeError(f"Expected a polars DataFrame/LazyFrame/Series, got {type(obj)}")
+    return np.asarray(obj)
 
 
 # ---------------------------------------------------------------------------
-# Public helpers
+# Legacy helper — kept for backward compatibility with external callers.
+# Internal learner code now uses filter_mask / filter_index / to_numpy.
 # ---------------------------------------------------------------------------
 
 
 def convert_pd_to_np(*args):
     """Convert pandas or polars objects to NumPy arrays.
 
-    Accepts any mix of:
-    * ``pd.DataFrame`` / ``pd.Series``   → ``.to_numpy()``
-    * ``pl.DataFrame`` / ``pl.LazyFrame`` / ``pl.Series``  → converted via
-      :func:`_polars_to_numpy`
-    * Any other type (e.g. ``np.ndarray``, ``None``) → returned unchanged.
+    .. deprecated::
+        Internal learner code no longer calls this function.  It is kept
+        solely for backward compatibility with user code that imports it
+        directly.  New code should use :func:`to_numpy` instead.
     """
-
-    def _convert(obj):
-        if obj is None:
-            return obj
-        if _POLARS_AVAILABLE and isinstance(
-            obj, (pl.DataFrame, pl.LazyFrame, pl.Series)
-        ):
-            return _polars_to_numpy(obj)
-        if hasattr(obj, "to_numpy"):
-            return obj.to_numpy()
-        return obj
-
-    output = [_convert(obj) for obj in args]
+    output = [to_numpy(obj) for obj in args]
     return output if len(output) > 1 else output[0]
 
 
+# ---------------------------------------------------------------------------
+# Validation helpers
+# ---------------------------------------------------------------------------
+
+
 def check_treatment_vector(treatment, control_name=None):
-    n_unique_treatments = np.unique(treatment).shape[0]
+    """Assert that *treatment* has at least two unique levels."""
+    # Normalise to numpy for np.unique
+    t = to_numpy(treatment) if not isinstance(treatment, np.ndarray) else treatment
+    n_unique_treatments = np.unique(t).shape[0]
     assert n_unique_treatments > 1, "Treatment vector must have at least two levels."
     if control_name is not None:
         assert (
-            control_name in treatment
+            control_name in t
         ), "Control group level {} not found in treatment vector.".format(control_name)
 
 
 def check_p_conditions(p, t_groups):
     eps = np.finfo(float).eps
 
-    # Build the allowed types tuple dynamically so it works whether or not
-    # polars is installed.
     _allowed = [np.ndarray, pd.Series]
     if _POLARS_AVAILABLE:
         _allowed.append(pl.Series)
@@ -106,7 +199,6 @@ def check_p_conditions(p, t_groups):
     ), "p must be an np.ndarray, pd.Series, pl.Series (if polars is installed), or dict type"
 
     if isinstance(p, _allowed_tuple):
-        # Normalise to numpy for the value checks below
         p_np = p.to_numpy() if hasattr(p, "to_numpy") else np.asarray(p)
         assert (
             t_groups.shape[0] == 1
@@ -144,20 +236,13 @@ def check_explain_conditions(method, models, X=None, treatment=None, y=None):
         ), "X, treatment, and y must be provided if method = {}".format(method)
 
 
-def clean_xgboost_objective(objective):
-    """
-    Translate objective to be compatible with loaded xgboost version
-
-    Args
-    ----
+# ---------------------------------------------------------------------------
+# XGBoost objective helpers (unchanged)
+# ---------------------------------------------------------------------------
 
-    objective : string
-        The objective to translate.
 
-    Returns
-    -------
-    The translated objective, or original if no translation was required.
-    """
+def clean_xgboost_objective(objective):
+    """Translate objective to be compatible with the loaded xgboost version."""
     compat_before_v83 = {"reg:squarederror": "reg:linear"}
     compat_v83_or_later = {"reg:linear": "reg:squarederror"}
     if version.parse(xgboost_version) < version.parse("0.83"):
@@ -170,19 +255,7 @@ def clean_xgboost_objective(objective):
 
 
 def get_xgboost_objective_metric(objective):
-    """
-    Get the xgboost version-compatible objective and evaluation metric from a potentially version-incompatible input.
-
-    Args
-    ----
-
-    objective : string
-        An xgboost objective that may be incompatible with the installed version.
-
-    Returns
-    -------
-    A tuple with the translated objective and evaluation metric.
-    """
+    """Return version-compatible (objective, eval_metric) tuple."""
 
     def clean_dict_keys(orig):
         return {clean_xgboost_objective(k): v for (k, v) in orig.items()}
@@ -190,9 +263,7 @@ def clean_dict_keys(orig):
     metric_mapping = clean_dict_keys(
         {"rank:pairwise": "auc", "reg:squarederror": "rmse"}
     )
-
     objective = clean_xgboost_objective(objective)
-
     assert (
         objective in metric_mapping
     ), "Effect learner objective must be one of: " + ", ".join(metric_mapping)
@@ -200,22 +271,7 @@ def clean_dict_keys(orig):
 
 
 def get_weighted_variance(x, sample_weight):
-    """
-    Calculate the variance of array x with sample_weight.
-
-    Args
-    ----
-
-    x : (np.array)
-        A list of number
-
-    sample_weight (np.array or list): an array of sample weights indicating the
-        weight of each observation for `effect_learner`. If None, it assumes equal weight.
-
-    Returns
-    -------
-    The variance of x with sample weight
-    """
+    """Calculate the variance of array x with sample_weight."""
     average = np.average(x, weights=sample_weight)
     variance = np.average((x - average) ** 2, weights=sample_weight)
     return variance
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index 88b5dc1d..b53e2c47 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -7,7 +7,8 @@
 from causalml.inference.meta.base import BaseLearner
 from causalml.inference.meta.utils import (
     check_treatment_vector,
-    convert_pd_to_np,
+    filter_mask,
+    to_numpy,
 )
 from causalml.metrics import regression_metrics, classification_metrics
 
@@ -15,12 +16,7 @@
 
 
 class BaseXLearner(BaseLearner):
-    """A parent class for X-learner regressor classes.
-
-    An X-learner estimates treatment effects with four machine learning models.
-
-    Details of X-learner are available at `Kunzel et al. (2018) <https://arxiv.org/abs/1706.03461>`_.
-    """
+    """A parent class for X-learner regressor classes."""
 
     def __init__(
         self,
@@ -32,18 +28,6 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
-        """Initialize a X-learner.
-
-        Args:
-            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
-                groups
-            control_outcome_learner (optional): a model to estimate outcomes in the control group
-            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
-            control_effect_learner (optional): a model to estimate treatment effects in the control group
-            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         assert (learner is not None) or (
             (control_outcome_learner is not None)
             and (treatment_outcome_learner is not None)
@@ -51,29 +35,29 @@ def __init__(
             and (treatment_effect_learner is not None)
         )
 
-        if control_outcome_learner is None:
-            self.model_mu_c = deepcopy(learner)
-        else:
-            self.model_mu_c = control_outcome_learner
-
-        if treatment_outcome_learner is None:
-            self.model_mu_t = deepcopy(learner)
-        else:
-            self.model_mu_t = treatment_outcome_learner
-
-        if control_effect_learner is None:
-            self.model_tau_c = deepcopy(learner)
-        else:
-            self.model_tau_c = control_effect_learner
-
-        if treatment_effect_learner is None:
-            self.model_tau_t = deepcopy(learner)
-        else:
-            self.model_tau_t = treatment_effect_learner
+        self.model_mu_c = (
+            deepcopy(learner)
+            if control_outcome_learner is None
+            else control_outcome_learner
+        )
+        self.model_mu_t = (
+            deepcopy(learner)
+            if treatment_outcome_learner is None
+            else treatment_outcome_learner
+        )
+        self.model_tau_c = (
+            deepcopy(learner)
+            if control_effect_learner is None
+            else control_effect_learner
+        )
+        self.model_tau_t = (
+            deepcopy(learner)
+            if treatment_effect_learner is None
+            else treatment_effect_learner
+        )
 
         self.ate_alpha = ate_alpha
         self.control_name = control_name
-
         self.propensity = None
         self.propensity_model = None
 
@@ -92,23 +76,16 @@ def __repr__(self):
         )
 
     def fit(self, X, treatment, y, p=None):
-        """Fit the inference model.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        treatment_np = to_numpy(treatment)
+        self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
 
         if p is None:
-            self._set_propensity_models(X=X, treatment=treatment, y=y)
+            # base.py does raw numpy indexing internally — convert at this boundary
+            self._set_propensity_models(
+                X=to_numpy(X), treatment=treatment_np, y=to_numpy(y)
+            )
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -126,69 +103,57 @@ def fit(self, X, treatment, y, p=None):
         self.vars_t = {}
 
         for group in self.t_groups:
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            X_filt = X[mask]
-            y_filt = y[mask]
-            w = (treatment_filt == group).astype(int)
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = filter_mask(treatment, mask)
+            X_filt = filter_mask(X, mask)
+            y_filt = filter_mask(y, mask)
+            w = (to_numpy(treatment_filt) == group).astype(int)
+
+            self.models_mu_c[group].fit(
+                filter_mask(X_filt, w == 0), filter_mask(y_filt, w == 0)
+            )
+            self.models_mu_t[group].fit(
+                filter_mask(X_filt, w == 1), filter_mask(y_filt, w == 1)
+            )
 
-            # Train outcome models
-            self.models_mu_c[group].fit(X_filt[w == 0], y_filt[w == 0])
-            self.models_mu_t[group].fit(X_filt[w == 1], y_filt[w == 1])
+            y_filt_np = to_numpy(y_filt)
+            X_filt_c = filter_mask(X_filt, w == 0)
+            X_filt_t = filter_mask(X_filt, w == 1)
 
-            # Calculate variances and treatment effects
             var_c = (
-                y_filt[w == 0] - self.models_mu_c[group].predict(X_filt[w == 0])
+                y_filt_np[w == 0] - self.models_mu_c[group].predict(X_filt_c)
             ).var()
             self.vars_c[group] = var_c
             var_t = (
-                y_filt[w == 1] - self.models_mu_t[group].predict(X_filt[w == 1])
+                y_filt_np[w == 1] - self.models_mu_t[group].predict(X_filt_t)
             ).var()
             self.vars_t[group] = var_t
 
-            # Train treatment models
-            d_c = self.models_mu_t[group].predict(X_filt[w == 0]) - y_filt[w == 0]
-            d_t = y_filt[w == 1] - self.models_mu_c[group].predict(X_filt[w == 1])
-            self.models_tau_c[group].fit(X_filt[w == 0], d_c)
-            self.models_tau_t[group].fit(X_filt[w == 1], d_t)
+            d_c = self.models_mu_t[group].predict(X_filt_c) - y_filt_np[w == 0]
+            d_t = y_filt_np[w == 1] - self.models_mu_c[group].predict(X_filt_t)
+            self.models_tau_c[group].fit(X_filt_c, d_c)
+            self.models_tau_t[group].fit(X_filt_t, d_t)
 
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        """Predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            return_components (bool, optional): whether to return outcome for treatment and control seperately
-            verbose (bool, optional): whether to output progress logs
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects.
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
-
         if p is None:
             logger.info("Generating propensity score")
-            p = dict()
-            for group in self.t_groups:
-                p_model = self.propensity_model[group]
-                p[group] = p_model.predict(X)
+            p = {
+                group: self.propensity_model[group].predict(X)
+                for group in self.t_groups
+            }
         else:
             p = self._format_p(p, self.t_groups)
 
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         dhat_cs = {}
         dhat_ts = {}
 
         for i, group in enumerate(self.t_groups):
-            model_tau_c = self.models_tau_c[group]
-            model_tau_t = self.models_tau_t[group]
-            dhat_cs[group] = model_tau_c.predict(X)
-            dhat_ts[group] = model_tau_t.predict(X)
+            dhat_cs[group] = self.models_tau_c[group].predict(X)
+            dhat_ts[group] = self.models_tau_t[group].predict(X)
 
             _te = (p[group] * dhat_cs[group] + (1 - p[group]) * dhat_ts[group]).reshape(
                 -1, 1
@@ -196,15 +161,20 @@ def predict(
             te[:, i] = np.ravel(_te)
 
             if (y is not None) and (treatment is not None) and verbose:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                X_filt = X[mask]
-                y_filt = y[mask]
-                w = (treatment_filt == group).astype(int)
+                treatment_np = to_numpy(treatment)
+                mask = (treatment_np == group) | (treatment_np == self.control_name)
+                treatment_filt_np = treatment_np[mask]
+                X_filt = filter_mask(X, mask)
+                y_filt = to_numpy(filter_mask(y, mask))
+                w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
-                yhat[w == 0] = self.models_mu_c[group].predict(X_filt[w == 0])
-                yhat[w == 1] = self.models_mu_t[group].predict(X_filt[w == 1])
+                yhat[w == 0] = self.models_mu_c[group].predict(
+                    filter_mask(X_filt, w == 0)
+                )
+                yhat[w == 1] = self.models_mu_t[group].predict(
+                    filter_mask(X_filt, w == 1)
+                )
 
                 logger.info("Error metrics for group {}".format(group))
                 regression_metrics(y_filt, yhat, w)
@@ -226,26 +196,6 @@ def fit_predict(
         return_components=False,
         verbose=True,
     ):
-        """Fit the treatment effect and outcome models of the R learner and predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            return_ci (bool): whether to return confidence intervals
-            n_bootstraps (int): number of bootstrap iterations
-            bootstrap_size (int): number of samples per bootstrap
-            return_components (bool, optional): whether to return outcome for treatment and control seperately
-            verbose (str): whether to output progress logs
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment]
-                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
-                UB [n_samples, n_treatment]
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         self.fit(X, treatment, y, p)
 
         if p is None:
@@ -260,6 +210,10 @@ def fit_predict(
         if not return_ci:
             return te
         else:
+            X_np = to_numpy(X)
+            treatment_np = to_numpy(treatment)
+            y_np = to_numpy(y)
+
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_mu_c_global = deepcopy(self.models_mu_c)
@@ -267,12 +221,12 @@ def fit_predict(
             models_tau_c_global = deepcopy(self.models_tau_c)
             models_tau_t_global = deepcopy(self.models_tau_t)
             te_bootstraps = np.zeros(
-                shape=(X.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
             )
 
             logger.info("Bootstrap Confidence Intervals")
             for i in tqdm(range(n_bootstraps)):
-                te_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
+                te_b = self.bootstrap(X_np, treatment_np, y_np, p, size=bootstrap_size)
                 te_bootstraps[:, :, i] = te_b
 
             te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
@@ -280,7 +234,6 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_mu_c = deepcopy(models_mu_c_global)
@@ -301,25 +254,8 @@ def estimate_ate(
         bootstrap_size=10000,
         pretrain=False,
     ):
-        """Estimate the Average Treatment Effect (ATE).
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            bootstrap_ci (bool): whether run bootstrap for confidence intervals
-            n_bootstraps (int): number of bootstrap iterations
-            bootstrap_size (int): number of samples per bootstrap
-            pretrain (bool): whether a model has been fit, default False.
-        Returns:
-            The mean and confidence interval (LB, UB) of the ATE estimate.
-        """
         if pretrain:
             if p is None:
-                # when p is null, use pretrain propensity score
                 if not self.propensity:
                     raise ValueError("no propensity score, please call fit() first")
                 te, dhat_cs, dhat_ts = self.predict(
@@ -334,7 +270,8 @@ def estimate_ate(
             te, dhat_cs, dhat_ts = self.fit_predict(
                 X, treatment, y, p, return_components=True
             )
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
+
+        treatment_np = to_numpy(treatment)
 
         if p is None:
             p = self.propensity
@@ -348,8 +285,8 @@ def estimate_ate(
         for i, group in enumerate(self.t_groups):
             _ate = te[:, i].mean()
 
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = treatment_np[mask]
             w = (treatment_filt == group).astype(int)
             prob_treatment = float(sum(w)) / w.shape[0]
 
@@ -357,8 +294,6 @@ def estimate_ate(
             dhat_t = dhat_ts[group][mask]
             p_filt = p[group][mask]
 
-            # SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.
-            # "Recent Developments in the Econometrics of Program Evaluation." Journal of Economic Literature
             se = np.sqrt(
                 (
                     self.vars_t[group] / prob_treatment
@@ -378,6 +313,7 @@ def estimate_ate(
         if not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
+            X_np = to_numpy(X)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_mu_c_global = deepcopy(self.models_mu_c)
@@ -389,7 +325,9 @@ def estimate_ate(
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
 
             for n in tqdm(range(n_bootstraps)):
-                cate_b = self.bootstrap(X, treatment, y, p, size=bootstrap_size)
+                cate_b = self.bootstrap(
+                    X_np, treatment_np, to_numpy(y), p, size=bootstrap_size
+                )
                 ate_bootstraps[:, n] = cate_b.mean(axis=0)
 
             ate_lower = np.percentile(
@@ -399,7 +337,6 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
-            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_mu_c = deepcopy(models_mu_c_global)
@@ -410,10 +347,6 @@ def estimate_ate(
 
 
 class BaseXRegressor(BaseXLearner):
-    """
-    A parent class for X-learner regressor classes.
-    """
-
     def __init__(
         self,
         learner=None,
@@ -424,18 +357,6 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
-        """Initialize an X-learner regressor.
-
-        Args:
-            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
-                groups
-            control_outcome_learner (optional): a model to estimate outcomes in the control group
-            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
-            control_effect_learner (optional): a model to estimate treatment effects in the control group
-            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         super().__init__(
             learner=learner,
             control_outcome_learner=control_outcome_learner,
@@ -448,10 +369,6 @@ def __init__(
 
 
 class BaseXClassifier(BaseXLearner):
-    """
-    A parent class for X-learner classifier classes.
-    """
-
     def __init__(
         self,
         outcome_learner=None,
@@ -463,24 +380,6 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
-        """Initialize an X-learner classifier.
-
-        Args:
-            outcome_learner (optional): a model to estimate outcomes in both the control and treatment groups.
-                Should be a classifier.
-            effect_learner (optional): a model to estimate treatment effects in both the control and treatment groups.
-                Should be a regressor.
-            control_outcome_learner (optional): a model to estimate outcomes in the control group.
-                Should be a classifier.
-            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group.
-                Should be a classifier.
-            control_effect_learner (optional): a model to estimate treatment effects in the control group.
-                Should be a regressor.
-            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
-                Should be a regressor.
-            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
-            control_name (str or int, optional): name of control group
-        """
         if outcome_learner is not None:
             control_outcome_learner = outcome_learner
             treatment_outcome_learner = outcome_learner
@@ -506,23 +405,16 @@ def __init__(
             )
 
     def fit(self, X, treatment, y, p=None):
-        """Fit the inference model.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
+        treatment_np = to_numpy(treatment)
+        self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
 
         if p is None:
-            self._set_propensity_models(X=X, treatment=treatment, y=y)
+            # base.py does raw numpy indexing internally — convert at this boundary
+            self._set_propensity_models(
+                X=to_numpy(X), treatment=treatment_np, y=to_numpy(y)
+            )
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -540,78 +432,65 @@ def fit(self, X, treatment, y, p=None):
         self.vars_t = {}
 
         for group in self.t_groups:
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            X_filt = X[mask]
-            y_filt = y[mask]
-            w = (treatment_filt == group).astype(int)
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt = filter_mask(treatment, mask)
+            X_filt = filter_mask(X, mask)
+            y_filt = filter_mask(y, mask)
+            w = (to_numpy(treatment_filt) == group).astype(int)
+
+            self.models_mu_c[group].fit(
+                filter_mask(X_filt, w == 0), filter_mask(y_filt, w == 0)
+            )
+            self.models_mu_t[group].fit(
+                filter_mask(X_filt, w == 1), filter_mask(y_filt, w == 1)
+            )
 
-            # Train outcome models
-            self.models_mu_c[group].fit(X_filt[w == 0], y_filt[w == 0])
-            self.models_mu_t[group].fit(X_filt[w == 1], y_filt[w == 1])
+            y_filt_np = to_numpy(y_filt)
+            X_filt_c = filter_mask(X_filt, w == 0)
+            X_filt_t = filter_mask(X_filt, w == 1)
 
-            # Calculate variances and treatment effects
             var_c = (
-                y_filt[w == 0]
-                - self.models_mu_c[group].predict_proba(X_filt[w == 0])[:, 1]
+                y_filt_np[w == 0]
+                - self.models_mu_c[group].predict_proba(X_filt_c)[:, 1]
             ).var()
             self.vars_c[group] = var_c
             var_t = (
-                y_filt[w == 1]
-                - self.models_mu_t[group].predict_proba(X_filt[w == 1])[:, 1]
+                y_filt_np[w == 1]
+                - self.models_mu_t[group].predict_proba(X_filt_t)[:, 1]
             ).var()
             self.vars_t[group] = var_t
 
-            # Train treatment models
             d_c = (
-                self.models_mu_t[group].predict_proba(X_filt[w == 0])[:, 1]
-                - y_filt[w == 0]
+                self.models_mu_t[group].predict_proba(X_filt_c)[:, 1]
+                - y_filt_np[w == 0]
             )
             d_t = (
-                y_filt[w == 1]
-                - self.models_mu_c[group].predict_proba(X_filt[w == 1])[:, 1]
+                y_filt_np[w == 1]
+                - self.models_mu_c[group].predict_proba(X_filt_t)[:, 1]
             )
-            self.models_tau_c[group].fit(X_filt[w == 0], d_c)
-            self.models_tau_t[group].fit(X_filt[w == 1], d_t)
+            self.models_tau_c[group].fit(X_filt_c, d_c)
+            self.models_tau_t[group].fit(X_filt_t, d_t)
 
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        """Predict treatment effects.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            return_components (bool, optional): whether to return outcome for treatment and control seperately
-            return_p_score (bool, optional): whether to return propensity score
-            verbose (bool, optional): whether to output progress logs
-        Returns:
-            (numpy.ndarray): Predictions of treatment effects.
-        """
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
-
         if p is None:
             logger.info("Generating propensity score")
-            p = dict()
-            for group in self.t_groups:
-                p_model = self.propensity_model[group]
-                p[group] = p_model.predict(X)
+            p = {
+                group: self.propensity_model[group].predict(X)
+                for group in self.t_groups
+            }
         else:
             p = self._format_p(p, self.t_groups)
 
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+        X_np = to_numpy(X)
+        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
         dhat_cs = {}
         dhat_ts = {}
 
         for i, group in enumerate(self.t_groups):
-            model_tau_c = self.models_tau_c[group]
-            model_tau_t = self.models_tau_t[group]
-            dhat_cs[group] = model_tau_c.predict(X)
-            dhat_ts[group] = model_tau_t.predict(X)
+            dhat_cs[group] = self.models_tau_c[group].predict(X)
+            dhat_ts[group] = self.models_tau_t[group].predict(X)
 
             _te = (p[group] * dhat_cs[group] + (1 - p[group]) * dhat_ts[group]).reshape(
                 -1, 1
@@ -619,19 +498,20 @@ def predict(
             te[:, i] = np.ravel(_te)
 
             if (y is not None) and (treatment is not None) and verbose:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                X_filt = X[mask]
-                y_filt = y[mask]
-                w = (treatment_filt == group).astype(int)
+                treatment_np = to_numpy(treatment)
+                mask = (treatment_np == group) | (treatment_np == self.control_name)
+                treatment_filt_np = treatment_np[mask]
+                X_filt = filter_mask(X, mask)
+                y_filt = to_numpy(filter_mask(y, mask))
+                w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
-                yhat[w == 0] = self.models_mu_c[group].predict_proba(X_filt[w == 0])[
-                    :, 1
-                ]
-                yhat[w == 1] = self.models_mu_t[group].predict_proba(X_filt[w == 1])[
-                    :, 1
-                ]
+                yhat[w == 0] = self.models_mu_c[group].predict_proba(
+                    filter_mask(X_filt, w == 0)
+                )[:, 1]
+                yhat[w == 1] = self.models_mu_t[group].predict_proba(
+                    filter_mask(X_filt, w == 1)
+                )[:, 1]
 
                 logger.info("Error metrics for group {}".format(group))
                 classification_metrics(y_filt, yhat, w)

From d7377436f61071532e7a58926fcf054ad332a90f Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Sun, 14 Jun 2026 00:55:39 +0530
Subject: [PATCH 04/15] refactor native pandas/Polars DataFrame support with
 LazyFrame collection across all meta-learners

---
 causalml/inference/meta/base.py      |  64 ++++---
 causalml/inference/meta/drlearner.py | 251 +++++++++++++++++++++----
 causalml/inference/meta/rlearner.py  | 267 ++++++++++++++++++++++-----
 causalml/inference/meta/slearner.py  | 143 ++++++++++++--
 causalml/inference/meta/tlearner.py  |  99 +++++-----
 causalml/inference/meta/utils.py     | 250 +++++++++++++++++++------
 causalml/inference/meta/xlearner.py  | 213 +++++++++++++++++----
 causalml/propensity.py               |  45 ++---
 pyproject.toml                       |   4 +-
 9 files changed, 1050 insertions(+), 286 deletions(-)

diff --git a/causalml/inference/meta/base.py b/causalml/inference/meta/base.py
index 50d7a019..faeb8bba 100644
--- a/causalml/inference/meta/base.py
+++ b/causalml/inference/meta/base.py
@@ -4,7 +4,13 @@
 import pandas as pd
 
 from causalml.inference.meta.explainer import Explainer
-from causalml.inference.meta.utils import check_p_conditions, convert_pd_to_np
+from causalml.inference.meta.utils import (
+    check_p_conditions,
+    filter_mask,
+    filter_index,
+    n_rows,
+    to_numpy,
+)
 from causalml.propensity import compute_propensity_score
 
 logger = logging.getLogger("causalml")
@@ -53,12 +59,27 @@ def estimate_ate(
         pass
 
     def bootstrap(self, X, treatment, y, p=None, size=10000, rng=None):
-        """Runs a single bootstrap. Fits on bootstrapped sample, then predicts on whole population."""
+        """Runs a single bootstrap. Fits on bootstrapped sample, then predicts on whole population.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, or pl.DataFrame): a feature matrix.
+                Resampled natively via :func:`filter_index`, so X stays in its
+                original format (numpy/pandas/polars) throughout.
+            treatment (np.array): a treatment vector (numpy)
+            y (np.array): an outcome vector (numpy)
+            p (dict, optional): a dict of {treatment group: propensity scores (numpy)}
+            size (int, optional): number of samples to draw with replacement
+            rng (np.random.Generator, optional): random number generator for
+                deterministic resampling
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects on the full X
+                from a model trained on the resampled subset.
+        """
         if rng is not None:
-            idxs = rng.choice(np.arange(0, X.shape[0]), size=size)
+            idxs = rng.choice(np.arange(0, n_rows(X)), size=size)
         else:
-            idxs = np.random.choice(np.arange(0, X.shape[0]), size=size)
-        X_b = X[idxs]
+            idxs = np.random.choice(np.arange(0, n_rows(X)), size=size)
+        X_b = filter_index(X, idxs)
 
         if p is not None:
             p_b = {group: _p[idxs] for group, _p in p.items()}
@@ -75,21 +96,19 @@ def _format_p(p, t_groups):
         """Format propensity scores into a dictionary of {treatment group: propensity scores}.
 
         Args:
-            p (np.ndarray, pd.Series, or dict): propensity scores
+            p (np.ndarray, pd.Series, pl.Series, or dict): propensity scores
             t_groups (list): treatment group names.
 
         Returns:
-            dict of {treatment group: propensity scores}
+            dict of {treatment group: propensity scores (numpy.ndarray)}
         """
         check_p_conditions(p, t_groups)
 
-        if isinstance(p, (np.ndarray, pd.Series)):
+        if isinstance(p, dict):
+            p = {treatment_name: to_numpy(_p) for treatment_name, _p in p.items()}
+        else:
             treatment_name = t_groups[0]
-            p = {treatment_name: convert_pd_to_np(p)}
-        elif isinstance(p, dict):
-            p = {
-                treatment_name: convert_pd_to_np(_p) for treatment_name, _p in p.items()
-            }
+            p = {treatment_name: to_numpy(p)}
 
         return p
 
@@ -103,19 +122,22 @@ def _set_propensity_models(self, X, treatment, y):
         PropensityModel (i.e. ElasticNetPropensityModel).
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
+            X (np.matrix, np.array, pd.DataFrame, or pl.DataFrame): a feature matrix.
+                Kept in its native format; scikit-learn >= 1.6 accepts pandas
+                and Polars DataFrames natively, so no conversion is performed.
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
         """
         logger.info("Generating propensity score")
+        treatment_np = to_numpy(treatment)
         p = dict()
         p_model = dict()
         for group in self.t_groups:
-            mask = (treatment == group) | (treatment == self.control_name)
-            treatment_filt = treatment[mask]
-            X_filt = X[mask]
-            w_filt = (treatment_filt == group).astype(int)
-            w = (treatment == group).astype(int)
+            mask = (treatment_np == group) | (treatment_np == self.control_name)
+            treatment_filt_np = treatment_np[mask]
+            X_filt = filter_mask(X, mask)
+            w_filt = (treatment_filt_np == group).astype(int)
+            w = (treatment_np == group).astype(int)
             propensity_model = self.model_p if hasattr(self, "model_p") else None
             p[group], p_model[group] = compute_propensity_score(
                 X=X_filt,
diff --git a/causalml/inference/meta/drlearner.py b/causalml/inference/meta/drlearner.py
index 58640c4e..ea928903 100644
--- a/causalml/inference/meta/drlearner.py
+++ b/causalml/inference/meta/drlearner.py
@@ -11,8 +11,10 @@
 from causalml.inference.meta.utils import (
     check_treatment_vector,
     check_p_conditions,
+    collect_if_lazy,
     filter_mask,
     filter_index,
+    n_rows,
     to_numpy,
 )
 from causalml.metrics import regression_metrics, classification_metrics
@@ -22,7 +24,12 @@
 
 
 class BaseDRLearner(BaseLearner):
-    """A parent class for DR-learner regressor classes."""
+    """A parent class for DR-learner regressor classes.
+
+    A DR-learner estimates treatment effects with machine learning models.
+
+    Details of DR-learner are available at `Kennedy (2020) <https://arxiv.org/abs/2004.14497>`_.
+    """
 
     def __init__(
         self,
@@ -33,6 +40,17 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
+        """Initialize a DR-learner.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups
+            control_outcome_learner (optional): a model to estimate outcomes in the control group
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         assert (learner is not None) or (
             (control_outcome_learner is not None)
             and (treatment_outcome_learner is not None)
@@ -57,6 +75,7 @@ def __init__(
 
         self.ate_alpha = ate_alpha
         self.control_name = control_name
+
         self.propensity = None
 
     def __repr__(self):
@@ -72,6 +91,21 @@ def __repr__(self):
         )
 
     def fit(self, X, treatment, y, p=None, seed=None):
+        """Fit the inference model.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method; the
+                feature matrix is otherwise kept in its native format throughout,
+                including the KFold partitions (sliced via :func:`filter_index`).
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            seed (int): random seed for cross-fitting
+        """
+        X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
@@ -80,19 +114,33 @@ def fit(self, X, treatment, y, p=None, seed=None):
         self.t_groups.sort()
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
 
+        # The estimator splits the data into 3 partitions for cross-fit on the propensity score estimation,
+        # the outcome regression, and the treatment regression on the doubly robust estimates. The use of
+        # the partitions is rotated so we do not lose on the sample size.
         cv = KFold(n_splits=3, shuffle=True, random_state=seed)
         split_indices = [index for _, index in cv.split(y_np)]
 
-        self.models_mu_c = [deepcopy(self.model_mu_c) for _ in range(3)]
+        self.models_mu_c = [
+            deepcopy(self.model_mu_c),
+            deepcopy(self.model_mu_c),
+            deepcopy(self.model_mu_c),
+        ]
         self.models_mu_t = {
-            group: [deepcopy(self.model_mu_t) for _ in range(3)]
+            group: [
+                deepcopy(self.model_mu_t),
+                deepcopy(self.model_mu_t),
+                deepcopy(self.model_mu_t),
+            ]
             for group in self.t_groups
         }
         self.models_tau = {
-            group: [deepcopy(self.model_tau) for _ in range(3)]
+            group: [
+                deepcopy(self.model_tau),
+                deepcopy(self.model_tau),
+                deepcopy(self.model_tau),
+            ]
             for group in self.t_groups
         }
-
         if p is None:
             self.propensity = {
                 group: np.zeros(y_np.shape[0]) for group in self.t_groups
@@ -104,12 +152,9 @@ def fit(self, X, treatment, y, p=None, seed=None):
             tau_idx = split_indices[(ifold + 2) % 3]
 
             treatment_treat = filter_index(treatment, treatment_idx)
-            treatment_out = filter_index(treatment, outcome_idx)
-            treatment_tau = filter_index(treatment, tau_idx)
-
-            treatment_treat_np = to_numpy(treatment_treat)
-            treatment_out_np = to_numpy(treatment_out)
-            treatment_tau_np = to_numpy(treatment_tau)
+            treatment_out_np = treatment_np[outcome_idx]
+            treatment_tau_np = treatment_np[tau_idx]
+            treatment_treat_np = treatment_np[treatment_idx]
 
             y_out = y_np[outcome_idx]
             y_tau = y_np[tau_idx]
@@ -121,13 +166,13 @@ def fit(self, X, treatment, y, p=None, seed=None):
             if p is None:
                 logger.info("Generating propensity score")
                 cur_p = dict()
+
                 for group in self.t_groups:
                     mask = (treatment_treat_np == group) | (
                         treatment_treat_np == self.control_name
                     )
-                    treatment_filt = filter_mask(treatment_treat, mask)
                     X_filt = filter_mask(X_treat, mask)
-                    w_filt = (to_numpy(treatment_filt) == group).astype(int)
+                    w_filt = (treatment_treat_np[mask] == group).astype(int)
                     w = (treatment_tau_np == group).astype(int)
                     cur_p[group], _ = compute_propensity_score(
                         X=X_filt, treatment=w_filt, X_pred=X_tau, treatment_pred=w
@@ -136,12 +181,9 @@ def fit(self, X, treatment, y, p=None, seed=None):
             else:
                 cur_p = dict()
                 if isinstance(p, (np.ndarray, pd.Series)):
-                    cur_p = {self.t_groups[0]: to_numpy(filter_index(p, tau_idx))}
+                    cur_p = {self.t_groups[0]: to_numpy(p)[tau_idx]}
                 else:
-                    cur_p = {
-                        g: to_numpy(filter_index(prop, tau_idx))
-                        for g, prop in p.items()
-                    }
+                    cur_p = {g: to_numpy(prop)[tau_idx] for g, prop in p.items()}
                 check_p_conditions(cur_p, self.t_groups)
 
             logger.info("Generate outcome regressions")
@@ -156,14 +198,14 @@ def fit(self, X, treatment, y, p=None, seed=None):
                 )
 
             logger.info("Fit pseudo outcomes from the DR formula")
+
             for group in self.t_groups:
                 mask = (treatment_tau_np == group) | (
                     treatment_tau_np == self.control_name
                 )
-                treatment_filt_np = treatment_tau_np[mask]
                 X_filt = filter_mask(X_tau, mask)
                 y_filt = y_tau[mask]
-                w_filt = (treatment_filt_np == group).astype(int)
+                w_filt = (treatment_tau_np[mask] == group).astype(int)
                 p_filt = cur_p[group][mask]
                 mu_t = self.models_mu_t[group][ifold].predict(X_filt)
                 mu_c = self.models_mu_c[ifold].predict(X_filt)
@@ -178,29 +220,62 @@ def fit(self, X, treatment, y, p=None, seed=None):
                 self.models_tau[group][ifold].fit(X_filt, dr)
 
     def bootstrap(self, X, treatment, y, p=None, size=10000, rng=None, seed=None):
+        """Runs a single bootstrap with optional deterministic cross-fit seed.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, or pl.DataFrame): a feature matrix.
+                Resampled natively via :func:`filter_index`.
+            treatment (np.array): a treatment vector (numpy)
+            y (np.array): an outcome vector (numpy)
+            p (dict, optional): a dict of {treatment group: propensity scores (numpy)}
+            size (int, optional): number of samples to draw with replacement
+            rng (np.random.Generator, optional): random number generator for
+                deterministic resampling
+            seed (int, optional): random seed for cross-fitting within the
+                resampled fit() call
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects on the full X
+                from a model trained on the resampled subset.
+        """
         if rng is not None:
-            idxs = rng.choice(np.arange(0, to_numpy(X).shape[0]), size=size)
+            idxs = rng.choice(np.arange(0, n_rows(X)), size=size)
         else:
-            idxs = np.random.choice(np.arange(0, to_numpy(X).shape[0]), size=size)
+            idxs = np.random.choice(np.arange(0, n_rows(X)), size=size)
         X_b = filter_index(X, idxs)
-        p_b = {group: _p[idxs] for group, _p in p.items()} if p is not None else None
-        treatment_b = filter_index(treatment, idxs)
-        y_b = to_numpy(y)[idxs]
+
+        if p is not None:
+            p_b = {group: _p[idxs] for group, _p in p.items()}
+        else:
+            p_b = None
+
+        treatment_b = treatment[idxs]
+        y_b = y[idxs]
         self.fit(X=X_b, treatment=treatment_b, y=y_b, p=p_b, seed=seed)
         return self.predict(X=X, p=p)
 
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
+            y (np.array, pd.Series, or pl.Series, optional): an outcome vector
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X = collect_if_lazy(X)
+
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         yhat_cs = {}
         yhat_ts = {}
 
         for i, group in enumerate(self.t_groups):
-            _te = np.r_[[model.predict(X) for model in self.models_tau[group]]].mean(
-                axis=0
-            )
+            models_tau = self.models_tau[group]
+            _te = np.r_[[model.predict(X) for model in models_tau]].mean(axis=0)
             te[:, i] = np.ravel(_te)
             yhat_cs[group] = np.r_[
                 [model.predict(X) for model in self.models_mu_c]
@@ -241,6 +316,27 @@ def fit_predict(
         verbose=True,
         seed=None,
     ):
+        """Fit the treatment effect and outcome models of the DR learner and predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            return_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (str): whether to output progress logs
+            seed (int): random seed for cross-fitting
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment]
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        X = collect_if_lazy(X)
         self.fit(X, treatment, y, p, seed)
 
         if p is None:
@@ -259,7 +355,6 @@ def fit_predict(
         if not return_ci:
             return te
         else:
-            X_np = to_numpy(X)
             treatment_np = to_numpy(treatment)
             y_np = to_numpy(y)
 
@@ -269,8 +364,9 @@ def fit_predict(
             models_mu_t_global = deepcopy(self.models_mu_t)
             models_tau_global = deepcopy(self.models_tau)
             te_bootstraps = np.zeros(
-                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(n_rows(X), self.t_groups.shape[0], n_bootstraps)
             )
+            # seed controls both bootstrap resampling and cross-fit randomness.
             rng = np.random.default_rng(seed) if seed is not None else None
 
             logger.info("Bootstrap Confidence Intervals")
@@ -281,7 +377,7 @@ def fit_predict(
                     else None
                 )
                 te_b = self.bootstrap(
-                    X_np,
+                    X,
                     treatment_np,
                     y_np,
                     p,
@@ -296,6 +392,7 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
+            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_mu_c = deepcopy(models_mu_c_global)
@@ -316,6 +413,25 @@ def estimate_ate(
         seed=None,
         pretrain=False,
     ):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            bootstrap_ci (bool): whether run bootstrap for confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            seed (int): random seed for cross-fitting
+            pretrain (bool): whether a model has been fit, default False.
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        X = collect_if_lazy(X)
+
         if pretrain:
             te, yhat_cs, yhat_ts = self.predict(
                 X, treatment, y, p, return_components=True
@@ -343,6 +459,7 @@ def estimate_ate(
 
         for i, group in enumerate(self.t_groups):
             _ate = te[:, i].mean()
+
             mask = (treatment_np == group) | (treatment_np == self.control_name)
             treatment_filt = treatment_np[mask]
             w = (treatment_filt == group).astype(int)
@@ -352,6 +469,8 @@ def estimate_ate(
             yhat_t = yhat_ts[group][mask]
             y_filt = y_np[mask]
 
+            # SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.
+            # "Recent Developments in the Econometrics of Program Evaluation." Journal of Economic Literature
             se = np.sqrt(
                 (
                     (y_filt[w == 0] - yhat_c[w == 0]).var() / (1 - prob_treatment)
@@ -371,7 +490,6 @@ def estimate_ate(
         if not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
-            X_np = to_numpy(X)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_mu_c_global = deepcopy(self.models_mu_c)
@@ -380,6 +498,7 @@ def estimate_ate(
 
             logger.info("Bootstrap Confidence Intervals for ATE")
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
+            # seed controls both bootstrap resampling and cross-fit randomness.
             rng = np.random.default_rng(seed) if seed is not None else None
 
             for n in tqdm(range(n_bootstraps)):
@@ -389,7 +508,7 @@ def estimate_ate(
                     else None
                 )
                 cate_b = self.bootstrap(
-                    X_np,
+                    X,
                     treatment_np,
                     y_np,
                     p,
@@ -406,6 +525,7 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
+            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_mu_c = deepcopy(models_mu_c_global)
@@ -415,6 +535,8 @@ def estimate_ate(
 
 
 class BaseDRRegressor(BaseDRLearner):
+    """A parent class for DR-learner regressor classes."""
+
     def __init__(
         self,
         learner=None,
@@ -424,6 +546,17 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
+        """Initialize an DR-learner regressor.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups
+            control_outcome_learner (optional): a model to estimate outcomes in the control group
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         super().__init__(
             learner=learner,
             control_outcome_learner=control_outcome_learner,
@@ -435,6 +568,8 @@ def __init__(
 
 
 class BaseDRClassifier(BaseDRLearner):
+    """A parent class for DR-learner classifier classes."""
+
     def __init__(
         self,
         learner=None,
@@ -444,6 +579,20 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
+        """Initialize a DR-learner classifier.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups. Should have a predict_proba() method for outcome models.
+            control_outcome_learner (optional): a model to estimate outcomes in the control group.
+                Should have a predict_proba() method.
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group.
+                Should have a predict_proba() method.
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group.
+                Should be a regressor.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         super().__init__(
             learner=learner,
             control_outcome_learner=control_outcome_learner,
@@ -456,15 +605,36 @@ def __init__(
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector. Used for computing
+                classification metrics when y is also provided.
+            y (np.array, pd.Series, or pl.Series, optional): an outcome vector. Used for computing
+                classification metrics when treatment is also provided.
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1). Currently not used in prediction but kept for API consistency.
+            return_components (bool, optional): whether to return outcome probabilities for treatment and control
+                groups separately. Defaults to False.
+            verbose (bool, optional): whether to output progress logs. Defaults to True.
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+            If return_components is True, also returns:
+                - dict: Predicted probabilities for the control group (yhat_cs).
+                - dict: Predicted probabilities for the treatment group (yhat_ts).
+        """
+        X = collect_if_lazy(X)
+
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         yhat_cs = {}
         yhat_ts = {}
 
         for i, group in enumerate(self.t_groups):
-            _te = np.r_[[model.predict(X) for model in self.models_tau[group]]].mean(
-                axis=0
-            )
+            models_tau = self.models_tau[group]
+            _te = np.r_[[model.predict(X) for model in models_tau]].mean(axis=0)
             te[:, i] = np.ravel(_te)
             yhat_cs[group] = np.r_[
                 [model.predict_proba(X)[:, 1] for model in self.models_mu_c]
@@ -495,6 +665,7 @@ def predict(
 
 class XGBDRRegressor(BaseDRRegressor):
     def __init__(self, ate_alpha=0.05, control_name=0, *args, **kwargs):
+        """Initialize a DR-learner with two XGBoost models."""
         super().__init__(
             learner=XGBRegressor(*args, **kwargs),
             ate_alpha=ate_alpha,
diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py
index 30b79d75..9e9dc422 100644
--- a/causalml/inference/meta/rlearner.py
+++ b/causalml/inference/meta/rlearner.py
@@ -9,7 +9,9 @@
 from causalml.inference.meta.base import BaseLearner
 from causalml.inference.meta.utils import (
     check_treatment_vector,
+    collect_if_lazy,
     filter_mask,
+    n_rows,
     to_numpy,
     get_xgboost_objective_metric,
     get_weighted_variance,
@@ -20,7 +22,12 @@
 
 
 class BaseRLearner(BaseLearner):
-    """A parent class for R-learner classes."""
+    """A parent class for R-learner classes.
+
+    An R-learner estimates treatment effects with two machine learning models and the propensity score.
+
+    Details of R-learner are available at `Nie and Wager (2019) <https://arxiv.org/abs/1712.04912>`_.
+    """
 
     def __init__(
         self,
@@ -34,6 +41,22 @@ def __init__(
         random_state=None,
         cv_n_jobs=-1,
     ):
+        """Initialize an R-learner.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects
+            outcome_learner (optional): a model to estimate outcomes
+            effect_learner (optional): a model to estimate treatment effects. It needs to take `sample_weight` as an
+                input argument for `fit()`
+            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
+                be used by default.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+            n_fold (int, optional): the number of cross validation folds for outcome_learner
+            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
+            cv_n_jobs (int, optional): number of parallel jobs to run for cross_val_predict. -1 means using all
+                processors
+        """
         assert (learner is not None) or (
             (outcome_learner is not None) and (effect_learner is not None)
         )
@@ -49,9 +72,11 @@ def __init__(
 
         self.ate_alpha = ate_alpha
         self.control_name = control_name
+
         self.random_state = random_state
         self.cv = KFold(n_splits=n_fold, shuffle=True, random_state=random_state)
         self.cv_n_jobs = cv_n_jobs
+
         self.propensity = None
         self.propensity_model = None
 
@@ -64,13 +89,31 @@ def __repr__(self):
         )
 
     def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
+        """Fit the treatment effect and outcome models of the R learner.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method; the
+                feature matrix is otherwise kept in its native format throughout,
+                including the call to ``cross_val_predict`` (scikit-learn >= 1.6
+                accepts pandas and Polars DataFrames natively).
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            verbose (bool, optional): whether to output progress logs
+        """
+        X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
 
         if sample_weight is not None:
             assert len(sample_weight) == len(
-                y
+                y_np
             ), "Data length must be equal for sample_weight and the input data"
             sample_weight = to_numpy(sample_weight)
 
@@ -78,7 +121,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
         self.t_groups.sort()
 
         if p is None:
-            self._set_propensity_models(X=to_numpy(X), treatment=treatment_np, y=y_np)
+            self._set_propensity_models(X=X, treatment=treatment_np, y=y_np)
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -90,7 +133,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
 
         if verbose:
             logger.info("generating out-of-fold CV outcome estimates")
-        # sklearn >= 1.6 accepts DataFrames natively — pass X as-is
         yhat = cross_val_predict(
             self.model_mu, X, y_np, cv=self.cv, n_jobs=self.cv_n_jobs
         )
@@ -131,8 +173,17 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             )
 
     def predict(self, X, p=None):
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X = collect_if_lazy(X)
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = self.models_tau[group].predict(X)
         return te
@@ -149,13 +200,33 @@ def fit_predict(
         bootstrap_size=10000,
         verbose=True,
     ):
+        """Fit the treatment effect and outcome models of the R learner and predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            return_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            verbose (bool): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        X = collect_if_lazy(X)
         self.fit(X, treatment, y, p, sample_weight, verbose=verbose)
         te = self.predict(X)
 
         if not return_ci:
             return te
         else:
-            X_np = to_numpy(X)
             treatment_np = to_numpy(treatment)
             y_np = to_numpy(y)
 
@@ -164,7 +235,7 @@ def fit_predict(
             model_mu_global = deepcopy(self.model_mu)
             models_tau_global = deepcopy(self.models_tau)
             te_bootstraps = np.zeros(
-                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(n_rows(X), self.t_groups.shape[0], n_bootstraps)
             )
 
             logger.info("Bootstrap Confidence Intervals")
@@ -173,7 +244,7 @@ def fit_predict(
                     p = self.propensity
                 else:
                     p = self._format_p(p, self.t_groups)
-                te_b = self.bootstrap(X_np, treatment_np, y_np, p, size=bootstrap_size)
+                te_b = self.bootstrap(X, treatment_np, y_np, p, size=bootstrap_size)
                 te_bootstraps[:, :, i] = te_b
 
             te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
@@ -181,6 +252,7 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
+            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.model_mu = deepcopy(model_mu_global)
@@ -200,13 +272,37 @@ def estimate_ate(
         bootstrap_size=10000,
         pretrain=False,
     ):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): only needed when pretrain=False, a treatment vector
+            y (np.array, pd.Series, or pl.Series): only needed when pretrain=False, an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            bootstrap_ci (bool): whether run bootstrap for confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            pretrain (bool): whether a model has been fit, default False.
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        X = collect_if_lazy(X)
         treatment_np = to_numpy(treatment)
-        X_np = to_numpy(X)
+        y_np = to_numpy(y)
 
         if pretrain:
             te = self.predict(X, p)
         else:
-            if not len(treatment_np) or not len(to_numpy(y)):
+            if (
+                treatment_np is None
+                or not len(treatment_np)
+                or y_np is None
+                or not len(y_np)
+            ):
                 raise ValueError("treatment and y must be provided when pretrain=False")
             te = self.fit_predict(X, treatment, y, p, sample_weight, return_ci=False)
 
@@ -216,17 +312,14 @@ def estimate_ate(
 
         for i, group in enumerate(self.t_groups):
             w = (treatment_np == group).astype(int)
-            prob_treatment = float(sum(w)) / X_np.shape[0]
+            prob_treatment = float(sum(w)) / n_rows(X)
             _ate = te[:, i].mean()
 
-            se = (
-                np.sqrt(
-                    (self.vars_t[group] / prob_treatment)
-                    + (self.vars_c[group] / (1 - prob_treatment))
-                    + te[:, i].var()
-                )
-                / X_np.shape[0]
-            )
+            se = np.sqrt(
+                (self.vars_t[group] / prob_treatment)
+                + (self.vars_c[group] / (1 - prob_treatment))
+                + te[:, i].var()
+            ) / n_rows(X)
 
             _ate_lb = _ate - se * norm.ppf(1 - self.ate_alpha / 2)
             _ate_ub = _ate + se * norm.ppf(1 - self.ate_alpha / 2)
@@ -238,7 +331,6 @@ def estimate_ate(
         if not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
-            y_np = to_numpy(y)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             model_mu_global = deepcopy(self.model_mu)
@@ -252,9 +344,7 @@ def estimate_ate(
                     p = self.propensity
                 else:
                     p = self._format_p(p, self.t_groups)
-                cate_b = self.bootstrap(
-                    X_np, treatment_np, y_np, p, size=bootstrap_size
-                )
+                cate_b = self.bootstrap(X, treatment_np, y_np, p, size=bootstrap_size)
                 ate_bootstraps[:, n] = cate_b.mean(axis=0)
 
             ate_lower = np.percentile(
@@ -264,6 +354,7 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
+            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.model_mu = deepcopy(model_mu_global)
@@ -272,6 +363,8 @@ def estimate_ate(
 
 
 class BaseRRegressor(BaseRLearner):
+    """A parent class for R-learner regressor classes."""
+
     def __init__(
         self,
         learner=None,
@@ -283,6 +376,20 @@ def __init__(
         n_fold=5,
         random_state=None,
     ):
+        """Initialize an R-learner regressor.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects
+            outcome_learner (optional): a model to estimate outcomes
+            effect_learner (optional): a model to estimate treatment effects. It needs to take `sample_weight` as an
+                input argument for `fit()`
+            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
+                be used by default.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+            n_fold (int, optional): the number of cross validation folds for outcome_learner
+            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
+        """
         super().__init__(
             learner=learner,
             outcome_learner=outcome_learner,
@@ -296,6 +403,8 @@ def __init__(
 
 
 class BaseRClassifier(BaseRLearner):
+    """A parent class for R-learner classifier classes."""
+
     def __init__(
         self,
         outcome_learner=None,
@@ -306,6 +415,19 @@ def __init__(
         n_fold=5,
         random_state=None,
     ):
+        """Initialize an R-learner classifier.
+
+        Args:
+            outcome_learner: a model to estimate outcomes. Should be a classifier.
+            effect_learner: a model to estimate treatment effects. It needs to take `sample_weight` as an
+                input argument for `fit()`. Should be a regressor.
+            propensity_learner (optional): a model to estimate propensity scores. `ElasticNetPropensityModel()` will
+                be used by default.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+            n_fold (int, optional): the number of cross validation folds for outcome_learner
+            random_state (int or RandomState, optional): a seed (int) or random number generator (RandomState)
+        """
         super().__init__(
             learner=None,
             outcome_learner=outcome_learner,
@@ -316,19 +438,35 @@ def __init__(
             n_fold=n_fold,
             random_state=random_state,
         )
+
         if (outcome_learner is None) and (effect_learner is None):
             raise ValueError(
                 "Either the outcome learner or the effect learner must be specified."
             )
 
     def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
+        """Fit the treatment effect and outcome models of the R learner.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            verbose (bool, optional): whether to output progress logs
+        """
+        X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
 
         if sample_weight is not None:
             assert len(sample_weight) == len(
-                y
+                y_np
             ), "Data length must be equal for sample_weight and the input data"
             sample_weight = to_numpy(sample_weight)
 
@@ -336,7 +474,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
         self.t_groups.sort()
 
         if p is None:
-            self._set_propensity_models(X=to_numpy(X), treatment=treatment_np, y=y_np)
+            self._set_propensity_models(X=X, treatment=treatment_np, y=y_np)
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -388,8 +526,17 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             )
 
     def predict(self, X, p=None):
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X = collect_if_lazy(X)
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = self.models_tau[group].predict(X)
         return te
@@ -407,7 +554,21 @@ def __init__(
         *args,
         **kwargs,
     ):
+        """Initialize an R-learner regressor with XGBoost model using pairwise ranking objective.
+
+        Args:
+            early_stopping: whether or not to use early stopping when fitting effect learner
+            test_size (float, optional): the proportion of the dataset to use as validation set when early stopping is
+                                         enabled
+            early_stopping_rounds (int, optional): validation metric needs to improve at least once in every
+                                                   early_stopping_rounds round(s) to continue training
+            effect_learner_objective (str, optional): the learning objective for the effect learner
+                                                      (default = 'reg:squarederror')
+            effect_learner_n_estimators (int, optional): number of trees to fit for the effect learner (default = 500)
+        """
+
         assert isinstance(random_state, int), "random_state should be int."
+
         objective, metric = get_xgboost_objective_metric(effect_learner_objective)
         self.effect_learner_objective = objective
         self.effect_learner_eval_metric = metric
@@ -416,6 +577,7 @@ def __init__(
         if self.early_stopping:
             self.test_size = test_size
             self.early_stopping_rounds = early_stopping_rounds
+
             effect_learner = XGBRegressor(
                 objective=self.effect_learner_objective,
                 n_estimators=self.effect_learner_n_estimators,
@@ -434,28 +596,43 @@ def __init__(
                 *args,
                 **kwargs,
             )
+
         super().__init__(
             outcome_learner=XGBRegressor(random_state=random_state, *args, **kwargs),
             effect_learner=effect_learner,
         )
 
     def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
+        """Fit the treatment effect and outcome models of the R learner.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
+                weight of each observation for `effect_learner`. If None, it assumes equal weight.
+            verbose (bool, optional): whether to output progress logs
+        """
+        X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
 
+        # initialize equal sample weight if it's not provided, for simplicity purpose
         sample_weight = (
             to_numpy(sample_weight) if sample_weight is not None else np.ones(len(y_np))
         )
         assert len(sample_weight) == len(
             y_np
         ), "Data length must be equal for sample_weight and the input data"
-
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
 
         if p is None:
-            self._set_propensity_models(X=to_numpy(X), treatment=treatment_np, y=y_np)
+            self._set_propensity_models(X=X, treatment=treatment_np, y=y_np)
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -470,15 +647,17 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
         yhat = cross_val_predict(self.model_mu, X, y_np, cv=self.cv, n_jobs=-1)
 
         for group in self.t_groups:
-            mask = (treatment_np == group) | (treatment_np == self.control_name)
-            treatment_filt_np = treatment_np[mask]
-            w = (treatment_filt_np == group).astype(int)
+            treatment_mask = (treatment_np == group) | (
+                treatment_np == self.control_name
+            )
+            treatment_filt = filter_mask(treatment, treatment_mask)
+            w = (to_numpy(treatment_filt) == group).astype(int)
 
-            X_filt = filter_mask(X, mask)
-            y_filt = y_np[mask]
-            yhat_filt = yhat[mask]
-            p_filt = p[group][mask]
-            sample_weight_filt = sample_weight[mask]
+            X_filt = filter_mask(X, treatment_mask)
+            y_filt = y_np[treatment_mask]
+            yhat_filt = yhat[treatment_mask]
+            p_filt = p[group][treatment_mask]
+            sample_weight_filt = sample_weight[treatment_mask]
 
             if verbose:
                 logger.info(
@@ -511,6 +690,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
                     test_size=self.test_size,
                     random_state=self.random_state,
                 )
+
                 self.models_tau[group].fit(
                     X=X_train_filt,
                     y=(y_train_filt - yhat_train_filt) / (w_train - p_train_filt),
@@ -527,6 +707,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
                     ],
                     verbose=verbose,
                 )
+
             else:
                 self.models_tau[group].fit(
                     X_filt,
@@ -536,9 +717,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
 
             diff_c = y_filt[w == 0] - yhat_filt[w == 0]
             diff_t = y_filt[w == 1] - yhat_filt[w == 1]
-            self.vars_c[group] = get_weighted_variance(
-                diff_c, sample_weight_filt[w == 0]
-            )
-            self.vars_t[group] = get_weighted_variance(
-                diff_t, sample_weight_filt[w == 1]
-            )
+            sample_weight_filt_c = sample_weight_filt[w == 0]
+            sample_weight_filt_t = sample_weight_filt[w == 1]
+            self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
+            self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
diff --git a/causalml/inference/meta/slearner.py b/causalml/inference/meta/slearner.py
index 445d2f4d..42040fc3 100644
--- a/causalml/inference/meta/slearner.py
+++ b/causalml/inference/meta/slearner.py
@@ -9,9 +9,11 @@
 from causalml.inference.meta.base import BaseLearner
 from causalml.inference.meta.utils import (
     check_treatment_vector,
+    collect_if_lazy,
+    concat_treatment_col,
     filter_mask,
+    n_rows,
     prepend_column,
-    concat_treatment_col,
     to_numpy,
 )
 from causalml.metrics import regression_metrics, classification_metrics
@@ -23,24 +25,50 @@ class StatsmodelsOLS:
     """A sklearn style wrapper class for statsmodels' OLS."""
 
     def __init__(self, cov_type="HC1", alpha=0.05):
+        """Initialize a statsmodels' OLS wrapper class object.
+
+        Args:
+            cov_type (str, optional): covariance estimator type.
+            alpha (float, optional): the confidence level alpha.
+        """
         self.cov_type = cov_type
         self.alpha = alpha
 
     def fit(self, X, y):
+        """Fit OLS.
+
+        Args:
+            X (np.matrix): a feature matrix
+            y (np.array): a label vector
+        """
+        # Append ones. The first column is for the treatment indicator.
         X = sm.add_constant(X, prepend=False, has_constant="add")
         self.model = sm.OLS(y, X).fit(cov_type=self.cov_type)
         self.coefficients = self.model.params
         self.conf_ints = self.model.conf_int(alpha=self.alpha)
 
     def predict(self, X):
+        # Append ones. The first column is for the treatment indicator.
         X = sm.add_constant(X, prepend=False, has_constant="add")
         return self.model.predict(X)
 
 
 class BaseSLearner(BaseLearner):
-    """A parent class for S-learner classes."""
+    """A parent class for S-learner classes.
+
+    An S-learner estimates treatment effects with one machine learning model.
+
+    Details of S-learner are available at `Kunzel et al. (2018) <https://arxiv.org/abs/1706.03461>`_.
+    """
 
     def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
+        """Initialize an S-learner.
+
+        Args:
+            learner (optional): a model to estimate the treatment effect
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         if learner is not None:
             self.model = learner
         else:
@@ -53,11 +81,15 @@ def __repr__(self):
 
     def fit(self, X, treatment, y, p=None):
         """Fit the inference model.
+
         Args:
-            X (np.matrix, np.array, pd.Dataframe, or pl.DataFrame): a feature matrix
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method; the
+                feature matrix is otherwise kept in its native format throughout.
             treatment (np.array, pd.Series, or pl.Series): a treatment vector
             y (np.array, pd.Series, or pl.Series): an outcome vector
         """
+        X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
@@ -79,8 +111,10 @@ def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
         """Predict treatment effects.
+
         Args:
-            X (np.matrix, np.array, pd.Dataframe, or pl.DataFrame): a feature matrix
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
             treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
             y (np.array, pd.Series, or pl.Series, optional): an outcome vector
             return_components (bool, optional): whether to return outcome for treatment and control seperately
@@ -88,12 +122,16 @@ def predict(
         Returns:
             (numpy.ndarray): Predictions of treatment effects.
         """
+        X = collect_if_lazy(X)
         yhat_cs = {}
         yhat_ts = {}
 
         for group in self.t_groups:
             model = self.models[group]
 
+            # Build separate frames for control and treatment to avoid in-place
+            # mutation, which fails when learners like CatBoost set the
+            # writeable flag to False on arrays passed to predict().
             X_new_c = prepend_column(0.0, X)
             yhat_cs[group] = model.predict(X_new_c)
 
@@ -114,8 +152,7 @@ def predict(
                 logger.info("Error metrics for group {}".format(group))
                 regression_metrics(y_filt, yhat, w)
 
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
@@ -136,13 +173,29 @@ def fit_predict(
         return_components=False,
         verbose=True,
     ):
+        """Fit the inference model of the S learner and predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            return_ci (bool, optional): whether to return confidence intervals
+            n_bootstraps (int, optional): number of bootstrap iterations
+            bootstrap_size (int, optional): number of samples per bootstrap
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment].
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        X = collect_if_lazy(X)
         self.fit(X, treatment, y)
         te = self.predict(X, treatment, y, return_components=return_components)
 
         if not return_ci:
             return te
         else:
-            X_np = to_numpy(X)
             treatment_np = to_numpy(treatment)
             y_np = to_numpy(y)
 
@@ -150,12 +203,12 @@ def fit_predict(
             _classes_global = self._classes
             models_global = deepcopy(self.models)
             te_bootstraps = np.zeros(
-                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(n_rows(X), self.t_groups.shape[0], n_bootstraps)
             )
 
             logger.info("Bootstrap Confidence Intervals")
             for i in tqdm(range(n_bootstraps)):
-                te_b = self.bootstrap(X_np, treatment_np, y_np, size=bootstrap_size)
+                te_b = self.bootstrap(X, treatment_np, y_np, size=bootstrap_size)
                 te_bootstraps[:, :, i] = te_b
 
             te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
@@ -163,6 +216,7 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
+            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models = deepcopy(models_global)
@@ -181,6 +235,21 @@ def estimate_ate(
         bootstrap_size=10000,
         pretrain=False,
     ):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            return_ci (bool, optional): whether to return confidence intervals
+            bootstrap_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            pretrain (bool): whether a model has been fit, default False.
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        X = collect_if_lazy(X)
         if pretrain:
             te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True)
         else:
@@ -228,7 +297,6 @@ def estimate_ate(
         elif return_ci and not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
-            X_np = to_numpy(X)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_global = deepcopy(self.models)
@@ -237,7 +305,7 @@ def estimate_ate(
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
 
             for n in tqdm(range(n_bootstraps)):
-                ate_b = self.bootstrap(X_np, treatment_np, y_np, size=bootstrap_size)
+                ate_b = self.bootstrap(X, treatment_np, y_np, size=bootstrap_size)
                 ate_bootstraps[:, n] = ate_b.mean(axis=0)
 
             ate_lower = np.percentile(
@@ -247,6 +315,7 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
+            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models = deepcopy(models_global)
@@ -255,14 +324,33 @@ def estimate_ate(
 
 
 class BaseSRegressor(BaseSLearner):
+    """A parent class for S-learner regressor classes."""
+
     def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
+        """Initialize an S-learner regressor.
+
+        Args:
+            learner (optional): a model to estimate the treatment effect
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         super().__init__(
             learner=learner, ate_alpha=ate_alpha, control_name=control_name
         )
 
 
 class BaseSClassifier(BaseSLearner):
+    """A parent class for S-learner classifier classes."""
+
     def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
+        """Initialize an S-learner classifier.
+
+        Args:
+            learner (optional): a model to estimate the treatment effect.
+                Should have a predict_proba() method.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         super().__init__(
             learner=learner, ate_alpha=ate_alpha, control_name=control_name
         )
@@ -270,6 +358,19 @@ def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
+            y (np.array, pd.Series, or pl.Series, optional): an outcome vector
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X = collect_if_lazy(X)
         yhat_cs = {}
         yhat_ts = {}
 
@@ -296,8 +397,7 @@ def predict(
                 logger.info("Error metrics for group {}".format(group))
                 classification_metrics(y_filt, yhat, w)
 
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
@@ -309,9 +409,26 @@ def predict(
 
 class LRSRegressor(BaseSRegressor):
     def __init__(self, ate_alpha=0.05, control_name=0):
+        """Initialize an S-learner with a linear regression model.
+
+        Args:
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         super().__init__(StatsmodelsOLS(alpha=ate_alpha), ate_alpha, control_name)
 
     def estimate_ate(self, X, treatment, y, p=None, pretrain=False):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            pretrain (bool): whether a model has been fit, default False.
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        X = collect_if_lazy(X)
         if not pretrain:
             self.fit(X, treatment, y)
 
diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py
index b80386d7..857cfa11 100644
--- a/causalml/inference/meta/tlearner.py
+++ b/causalml/inference/meta/tlearner.py
@@ -16,9 +16,10 @@
 
 from causalml.inference.meta.base import BaseLearner
 from causalml.inference.meta.utils import (
-    _POLARS_AVAILABLE,
     check_treatment_vector,
+    collect_if_lazy,
     filter_mask,
+    n_rows,
     to_numpy,
 )
 from causalml.metrics import regression_metrics, classification_metrics
@@ -75,23 +76,24 @@ def __repr__(self):
 
     @ignore_warnings(category=ConvergenceWarning)
     def fit(self, X, treatment, y, p=None):
-        """Fit the inference model
+        """Fit the inference model.
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
-            treatment (np.array or pd.Series or pl.Series): a treatment vector
-            y (np.array or pd.Series or pl.Series): an outcome vector
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method; the
+                feature matrix is otherwise kept in its native format throughout.
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
         """
+        X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
-        self.t_groups = np.unique(
-            to_numpy(treatment)[to_numpy(treatment) != self.control_name]
-        )
+        treatment_np = to_numpy(treatment)
+        self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
         self.models_c = {group: deepcopy(self.model_c) for group in self.t_groups}
         self.models_t = {group: deepcopy(self.model_t) for group in self.t_groups}
 
-        treatment_np = to_numpy(treatment)
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
             treatment_filt = filter_mask(treatment, mask)
@@ -112,21 +114,16 @@ def predict(
         """Predict treatment effects.
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
-            treatment (np.array or pd.Series or pl.Series, optional): a treatment vector
-            y (np.array or pd.Series or pl.Series, optional): an outcome vector
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
+            y (np.array, pd.Series, or pl.Series, optional): an outcome vector
             return_components (bool, optional): whether to return outcome for treatment and control seperately
             verbose (bool, optional): whether to output progress logs
         Returns:
             (numpy.ndarray): Predictions of treatment effects.
         """
-
-        if _POLARS_AVAILABLE:
-            import polars as pl
-
-            if isinstance(X, pl.LazyFrame):
-                X = X.collect()
-
+        X = collect_if_lazy(X)
         yhat_cs = {}
         yhat_ts = {}
 
@@ -150,8 +147,7 @@ def predict(
                 logger.info("Error metrics for group {}".format(group))
                 regression_metrics(y_filt, yhat, w)
 
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
@@ -175,9 +171,9 @@ def fit_predict(
         """Fit the inference model of the T learner and predict treatment effects.
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
-            treatment (np.array or pd.Series or pl.Series): a treatment vector
-            y (np.array or pd.Series or pl.Series): an outcome vector
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
             return_ci (bool): whether to return confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -188,13 +184,13 @@ def fit_predict(
                 If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
                 UB [n_samples, n_treatment]
         """
+        X = collect_if_lazy(X)
         self.fit(X, treatment, y)
         te = self.predict(X, treatment, y, return_components=return_components)
 
         if not return_ci:
             return te
         else:
-            X_np = to_numpy(X)
             treatment_np = to_numpy(treatment)
             y_np = to_numpy(y)
 
@@ -203,12 +199,12 @@ def fit_predict(
             models_c_global = deepcopy(self.models_c)
             models_t_global = deepcopy(self.models_t)
             te_bootstraps = np.zeros(
-                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(n_rows(X), self.t_groups.shape[0], n_bootstraps)
             )
 
             logger.info("Bootstrap Confidence Intervals")
             for i in tqdm(range(n_bootstraps)):
-                te_b = self.bootstrap(X_np, treatment_np, y_np, size=bootstrap_size)
+                te_b = self.bootstrap(X, treatment_np, y_np, size=bootstrap_size)
                 te_bootstraps[:, :, i] = te_b
 
             te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
@@ -216,6 +212,7 @@ def fit_predict(
                 te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2
             )
 
+            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_c = deepcopy(models_c_global)
@@ -237,9 +234,9 @@ def estimate_ate(
         """Estimate the Average Treatment Effect (ATE).
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
-            treatment (np.array or pd.Series or pl.Series): a treatment vector
-            y (np.array or pd.Series or pl.Series): an outcome vector
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
             bootstrap_ci (bool): whether to return confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -247,6 +244,7 @@ def estimate_ate(
         Returns:
             The mean and confidence interval (LB, UB) of the ATE estimate.
         """
+        X = collect_if_lazy(X)
         if pretrain:
             te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True)
         else:
@@ -254,7 +252,6 @@ def estimate_ate(
                 X, treatment, y, return_components=True
             )
 
-        # work in numpy for the ATE calculation
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
 
@@ -293,7 +290,6 @@ def estimate_ate(
         if not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
-            X_np = to_numpy(X)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_c_global = deepcopy(self.models_c)
@@ -303,7 +299,7 @@ def estimate_ate(
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
 
             for n in tqdm(range(n_bootstraps)):
-                ate_b = self.bootstrap(X_np, treatment_np, y_np, size=bootstrap_size)
+                ate_b = self.bootstrap(X, treatment_np, y_np, size=bootstrap_size)
                 ate_bootstraps[:, n] = ate_b.mean(axis=0)
 
             ate_lower = np.percentile(
@@ -313,6 +309,7 @@ def estimate_ate(
                 ate_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=1
             )
 
+            # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
             self.models_c = deepcopy(models_c_global)
@@ -332,6 +329,15 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
+        """Initialize a T-learner regressor.
+
+        Args:
+            learner (model): a model to estimate control and treatment outcomes.
+            control_learner (model, optional): a model to estimate control outcomes
+            treatment_learner (model, optional): a model to estimate treatment outcomes
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         super().__init__(
             learner=learner,
             control_learner=control_learner,
@@ -352,6 +358,15 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
+        """Initialize a T-learner classifier.
+
+        Args:
+            learner (model): a model to estimate control and treatment outcomes.
+            control_learner (model, optional): a model to estimate control outcomes
+            treatment_learner (model, optional): a model to estimate treatment outcomes
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         super().__init__(
             learner=learner,
             control_learner=control_learner,
@@ -366,19 +381,16 @@ def predict(
         """Predict treatment effects.
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe or pl.DataFrame): a feature matrix
-            treatment (np.array or pd.Series or pl.Series, optional): a treatment vector
-            y (np.array or pd.Series or pl.Series, optional): an outcome vector
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
+            y (np.array, pd.Series, or pl.Series, optional): an outcome vector
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
             verbose (bool, optional): whether to output progress logs
         Returns:
             (numpy.ndarray): Predictions of treatment effects.
         """
-        if _POLARS_AVAILABLE:
-            import polars as pl
-
-            if isinstance(X, pl.LazyFrame):
-                X = X.collect()
-
+        X = collect_if_lazy(X)
         yhat_cs = {}
         yhat_ts = {}
 
@@ -402,8 +414,7 @@ def predict(
                 logger.info("Error metrics for group {}".format(group))
                 classification_metrics(y_filt, yhat, w)
 
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = yhat_ts[group] - yhat_cs[group]
 
diff --git a/causalml/inference/meta/utils.py b/causalml/inference/meta/utils.py
index 9aeddb85..14b39b7c 100644
--- a/causalml/inference/meta/utils.py
+++ b/causalml/inference/meta/utils.py
@@ -18,35 +18,88 @@
 
 # ---------------------------------------------------------------------------
 # Native DataFrame helpers
+#
+# Design contract (per maintainer review):
+#   - X (feature matrices) stay native end-to-end: numpy, pandas, or polars.
+#     pl.LazyFrame is collected ONCE at the top of each public method into a
+#     pl.DataFrame via `collect_if_lazy`. After that point, helpers only need
+#     to handle pl.DataFrame / pl.Series.
+#   - treatment / y / p / sample_weight (1-D vectors used for masking,
+#     np.unique, .astype, etc.) are normalized to numpy via `to_numpy` at the
+#     entry of each public method.
+#   - Never call to_numpy(X) merely to read a row count: use `n_rows(X)`,
+#     which works for numpy, pandas, and polars.
 # ---------------------------------------------------------------------------
 
 
+def collect_if_lazy(X):
+    """Collect a polars LazyFrame into a DataFrame; pass through otherwise.
+
+    Call this once at the top of each public method on the feature matrix X
+    so that downstream helpers (filter_mask, filter_index, prepend_column,
+    concat_treatment_col, n_rows) never need to handle pl.LazyFrame.
+
+    Args:
+        X: numpy array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame.
+
+    Returns:
+        X unchanged, unless it was a pl.LazyFrame (then its .collect()).
+    """
+    if _POLARS_AVAILABLE and isinstance(X, pl.LazyFrame):
+        return X.collect()
+    return X
+
+
+def n_rows(X):
+    """Return the number of rows of X without converting to numpy.
+
+    Works for numpy arrays, pandas DataFrames/Series, and polars
+    DataFrames/Series (and LazyFrames, which are collected first).
+
+    Args:
+        X: numpy array, pd.DataFrame, pd.Series, pl.DataFrame, pl.Series,
+           or pl.LazyFrame.
+
+    Returns:
+        int: number of rows.
+    """
+    X = collect_if_lazy(X)
+    if hasattr(X, "shape"):
+        return X.shape[0]
+    return len(X)
+
+
 def filter_mask(obj, mask):
     """Filter rows by a boolean mask.
 
     Works natively for numpy arrays, pandas DataFrames/Series, and
-    polars DataFrames/Series without materialising unnecessary copies.
+    polars DataFrames/Series/LazyFrames without materialising unnecessary
+    copies.
 
     Args:
-        obj: numpy array, pd.DataFrame, pd.Series, pl.DataFrame, or pl.Series.
+        obj: numpy array, pd.DataFrame, pd.Series, pl.DataFrame, pl.Series,
+             or pl.LazyFrame.
         mask: boolean array or Series of the same length as obj.
 
     Returns:
-        Filtered object of the same type as input.
+        Filtered object of the same type as input (a pl.LazyFrame input is
+        returned as a pl.DataFrame, since filtering requires collection).
     """
     if obj is None:
         return None
+
+    obj = collect_if_lazy(obj)
+
     if isinstance(obj, (pd.DataFrame, pd.Series)):
         if isinstance(mask, pd.Series):
             return obj.loc[mask]
-        # numpy boolean array — reset index to avoid alignment issues
         return obj.loc[np.asarray(mask, dtype=bool)]
-    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.LazyFrame, pl.Series)):
-        if isinstance(obj, pl.LazyFrame):
-            obj = obj.collect()
+
+    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.Series)):
         if isinstance(mask, pl.Series):
             return obj.filter(mask)
         return obj.filter(pl.Series(np.asarray(mask, dtype=bool)))
+
     # numpy / anything else
     return obj[np.asarray(mask, dtype=bool)]
 
@@ -54,45 +107,58 @@ def filter_mask(obj, mask):
 def filter_index(obj, indices):
     """Filter rows by integer positional indices.
 
-    Used for KFold partition slicing (e.g. in DR-learner).
+    Used for KFold partition slicing (e.g. in DR-learner) and bootstrap
+    resampling.
 
     Args:
-        obj: numpy array, pd.DataFrame, pd.Series, pl.DataFrame, or pl.Series.
+        obj: numpy array, pd.DataFrame, pd.Series, pl.DataFrame, pl.Series,
+             or pl.LazyFrame.
         indices: integer array of row positions.
 
     Returns:
-        Filtered object of the same type as input.
+        Filtered object of the same type as input (a pl.LazyFrame input is
+        returned as a pl.DataFrame).
     """
     if obj is None:
         return None
+
+    obj = collect_if_lazy(obj)
+
     if isinstance(obj, (pd.DataFrame, pd.Series)):
         return obj.iloc[indices]
-    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.LazyFrame, pl.Series)):
-        if isinstance(obj, pl.LazyFrame):
-            obj = obj.collect()
+
+    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.Series)):
         return obj[indices]
+
     return obj[indices]  # numpy
 
 
 def prepend_column(value, X):
     """Prepend a constant-value column to a feature matrix X.
 
-    Used by the S-learner to prepend the treatment indicator before fitting.
+    Used by the S-learner to prepend the treatment indicator before
+    calling predict() for the control (value=0.0) and treatment (value=1.0)
+    counterfactuals.
 
     Args:
         value (float): scalar value to fill the new column with (0.0 or 1.0).
-        X: numpy array, pd.DataFrame, or pl.DataFrame.
+        X: numpy array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame.
 
     Returns:
-        Feature matrix with the new column prepended, same type as X.
+        Feature matrix with the new column prepended, same type as X (a
+        pl.LazyFrame input is returned as a pl.DataFrame).
     """
-    n = len(X)
+    X = collect_if_lazy(X)
+    n = n_rows(X)
+
     if isinstance(X, pd.DataFrame):
         col = pd.DataFrame({"_w": np.full(n, value)}, index=X.index)
         return pd.concat([col, X], axis=1)
+
     if _POLARS_AVAILABLE and isinstance(X, pl.DataFrame):
         col = pl.Series("_w", np.full(n, value))
         return X.with_columns(col).select(["_w"] + X.columns)
+
     # numpy
     return np.hstack((np.full((n, 1), value), X))
 
@@ -101,34 +167,39 @@ def concat_treatment_col(w, X):
     """Prepend an array-valued treatment column *w* to feature matrix X.
 
     Unlike :func:`prepend_column`, this takes an existing 1-D array rather
-    than a scalar. Used by S-learner ``fit()`` to build the augmented matrix.
+    than a scalar. Used by S-learner ``fit()`` to build the augmented
+    matrix [w | X].
 
     Args:
         w: 1-D numpy array of treatment indicators (0/1).
-        X: numpy array, pd.DataFrame, or pl.DataFrame.
+        X: numpy array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame.
 
     Returns:
-        Augmented feature matrix with w prepended, same type as X.
+        Augmented feature matrix with w prepended, same type as X (a
+        pl.LazyFrame input is returned as a pl.DataFrame).
     """
+    X = collect_if_lazy(X)
+
     if isinstance(X, pd.DataFrame):
-        col = pd.DataFrame({"_w": w}, index=X.index)
-        return pd.concat(
-            [col, X.reset_index(drop=True) if not X.index.equals(col.index) else X],
-            axis=1,
-        )
+        col = pd.DataFrame({"_w": np.asarray(w)}, index=X.index)
+        return pd.concat([col, X], axis=1)
+
     if _POLARS_AVAILABLE and isinstance(X, pl.DataFrame):
         col = pl.Series("_w", np.asarray(w))
         return X.with_columns(col).select(["_w"] + X.columns)
+
     # numpy
     return np.hstack((np.asarray(w).reshape(-1, 1), X))
 
 
 def to_numpy(obj):
-    """Convert a pandas or polars object to a NumPy array.
+    """Convert a pandas or polars 1-D vector (or feature matrix) to NumPy.
 
-    Use this *only* at boundaries where a third-party library strictly
-    requires numpy (rare — sklearn >= 1.6 and XGBoost >= 3.1 both accept
-    DataFrames natively).
+    Per the native-X contract, this should be called on ``treatment``,
+    ``y``, ``p``, and ``sample_weight`` at the entry of each public method.
+    It should NOT be called on the feature matrix X except at the small
+    number of boundaries that strictly require numpy (documented inline
+    where used).
 
     Args:
         obj: numpy array, pd.DataFrame, pd.Series, pl.DataFrame,
@@ -139,34 +210,19 @@ def to_numpy(obj):
     """
     if obj is None:
         return None
-    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.LazyFrame, pl.Series)):
-        if isinstance(obj, pl.LazyFrame):
-            obj = obj.collect()
+
+    obj = collect_if_lazy(obj)
+
+    if _POLARS_AVAILABLE and isinstance(obj, (pl.DataFrame, pl.Series)):
         arr = obj.to_numpy()
         if isinstance(obj, pl.DataFrame) and arr.ndim == 2 and arr.shape[1] == 1:
             return arr.ravel()
         return arr
+
     if hasattr(obj, "to_numpy"):
         return obj.to_numpy()
-    return np.asarray(obj)
 
-
-# ---------------------------------------------------------------------------
-# Legacy helper — kept for backward compatibility with external callers.
-# Internal learner code now uses filter_mask / filter_index / to_numpy.
-# ---------------------------------------------------------------------------
-
-
-def convert_pd_to_np(*args):
-    """Convert pandas or polars objects to NumPy arrays.
-
-    .. deprecated::
-        Internal learner code no longer calls this function.  It is kept
-        solely for backward compatibility with user code that imports it
-        directly.  New code should use :func:`to_numpy` instead.
-    """
-    output = [to_numpy(obj) for obj in args]
-    return output if len(output) > 1 else output[0]
+    return np.asarray(obj)
 
 
 # ---------------------------------------------------------------------------
@@ -175,9 +231,13 @@ def convert_pd_to_np(*args):
 
 
 def check_treatment_vector(treatment, control_name=None):
-    """Assert that *treatment* has at least two unique levels."""
-    # Normalise to numpy for np.unique
-    t = to_numpy(treatment) if not isinstance(treatment, np.ndarray) else treatment
+    """Assert that *treatment* has at least two unique levels.
+
+    Args:
+        treatment (np.array, pd.Series, or pl.Series): a treatment vector.
+        control_name (str or int, optional): name of the control group.
+    """
+    t = to_numpy(treatment)
     n_unique_treatments = np.unique(t).shape[0]
     assert n_unique_treatments > 1, "Treatment vector must have at least two levels."
     if control_name is not None:
@@ -187,6 +247,12 @@ def check_treatment_vector(treatment, control_name=None):
 
 
 def check_p_conditions(p, t_groups):
+    """Validate that propensity scores p lie strictly within (0, 1).
+
+    Args:
+        p (np.ndarray, pd.Series, pl.Series, or dict): propensity scores.
+        t_groups (np.ndarray): unique non-control treatment group labels.
+    """
     eps = np.finfo(float).eps
 
     _allowed = [np.ndarray, pd.Series]
@@ -199,7 +265,7 @@ def check_p_conditions(p, t_groups):
     ), "p must be an np.ndarray, pd.Series, pl.Series (if polars is installed), or dict type"
 
     if isinstance(p, _allowed_tuple):
-        p_np = p.to_numpy() if hasattr(p, "to_numpy") else np.asarray(p)
+        p_np = to_numpy(p)
         assert (
             t_groups.shape[0] == 1
         ), "If p is passed as an array/Series, there must be only 1 unique non-control group in the treatment vector."
@@ -209,14 +275,21 @@ def check_p_conditions(p, t_groups):
 
     if isinstance(p, dict):
         for t_name in t_groups:
-            p_val = p[t_name]
-            p_np = p_val.to_numpy() if hasattr(p_val, "to_numpy") else np.asarray(p_val)
+            p_np = to_numpy(p[t_name])
             assert (0 + eps < p_np).all() and (
                 p_np < 1 - eps
             ).all(), "The values of p should lie within the (0, 1) interval."
 
 
 def check_explain_conditions(method, models, X=None, treatment=None, y=None):
+    """Validate inputs for explain_validation-style feature-importance methods.
+
+    Args:
+        method (str): one of "gini", "permutation", "shapley".
+        models (list): models that need a ``.feature_importances_`` attribute
+            for "gini"/"shapley".
+        X, treatment, y: required (non-None) for "permutation"/"shapley".
+    """
     valid_methods = ["gini", "permutation", "shapley"]
     assert method in valid_methods, "Current supported methods: {}".format(
         ", ".join(valid_methods)
@@ -242,7 +315,19 @@ def check_explain_conditions(method, models, X=None, treatment=None, y=None):
 
 
 def clean_xgboost_objective(objective):
-    """Translate objective to be compatible with the loaded xgboost version."""
+    """
+    Translate objective to be compatible with loaded xgboost version
+
+    Args
+    ----
+
+    objective : string
+        The objective to translate.
+
+    Returns
+    -------
+    The translated objective, or original if no translation was required.
+    """
     compat_before_v83 = {"reg:squarederror": "reg:linear"}
     compat_v83_or_later = {"reg:linear": "reg:squarederror"}
     if version.parse(xgboost_version) < version.parse("0.83"):
@@ -255,7 +340,19 @@ def clean_xgboost_objective(objective):
 
 
 def get_xgboost_objective_metric(objective):
-    """Return version-compatible (objective, eval_metric) tuple."""
+    """
+    Get the xgboost version-compatible objective and evaluation metric from a potentially version-incompatible input.
+
+    Args
+    ----
+
+    objective : string
+        An xgboost objective that may be incompatible with the installed version.
+
+    Returns
+    -------
+    A tuple with the translated objective and evaluation metric.
+    """
 
     def clean_dict_keys(orig):
         return {clean_xgboost_objective(k): v for (k, v) in orig.items()}
@@ -263,7 +360,9 @@ def clean_dict_keys(orig):
     metric_mapping = clean_dict_keys(
         {"rank:pairwise": "auc", "reg:squarederror": "rmse"}
     )
+
     objective = clean_xgboost_objective(objective)
+
     assert (
         objective in metric_mapping
     ), "Effect learner objective must be one of: " + ", ".join(metric_mapping)
@@ -271,7 +370,40 @@ def clean_dict_keys(orig):
 
 
 def get_weighted_variance(x, sample_weight):
-    """Calculate the variance of array x with sample_weight."""
+    """
+    Calculate the variance of array x with sample_weight.
+
+    Args
+    ----
+
+    x : (np.array)
+        A list of number
+
+    sample_weight (np.array or list): an array of sample weights indicating the
+        weight of each observation for `effect_learner`. If None, it assumes equal weight.
+
+    Returns
+    -------
+    The variance of x with sample weight
+    """
     average = np.average(x, weights=sample_weight)
     variance = np.average((x - average) ** 2, weights=sample_weight)
     return variance
+
+
+# ---------------------------------------------------------------------------
+# Backward-compatibility alias
+# ---------------------------------------------------------------------------
+
+
+def convert_pd_to_np(*args):
+    """Deprecated alias for :func:`to_numpy`, kept for backward compatibility
+    with modules (e.g. explainer.py) that have not yet migrated to the
+    native-DataFrame contract.
+
+    .. deprecated::
+        Use :func:`to_numpy` for 1-D vectors. Do not use on feature matrices
+        X under the native-X contract.
+    """
+    output = [to_numpy(obj) for obj in args]
+    return output if len(output) > 1 else output[0]
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index b53e2c47..f12e886d 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -7,7 +7,9 @@
 from causalml.inference.meta.base import BaseLearner
 from causalml.inference.meta.utils import (
     check_treatment_vector,
+    collect_if_lazy,
     filter_mask,
+    n_rows,
     to_numpy,
 )
 from causalml.metrics import regression_metrics, classification_metrics
@@ -16,7 +18,12 @@
 
 
 class BaseXLearner(BaseLearner):
-    """A parent class for X-learner regressor classes."""
+    """A parent class for X-learner regressor classes.
+
+    An X-learner estimates treatment effects with four machine learning models.
+
+    Details of X-learner are available at `Kunzel et al. (2018) <https://arxiv.org/abs/1706.03461>`_.
+    """
 
     def __init__(
         self,
@@ -28,6 +35,18 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
+        """Initialize a X-learner.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups
+            control_outcome_learner (optional): a model to estimate outcomes in the control group
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
+            control_effect_learner (optional): a model to estimate treatment effects in the control group
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         assert (learner is not None) or (
             (control_outcome_learner is not None)
             and (treatment_outcome_learner is not None)
@@ -58,6 +77,7 @@ def __init__(
 
         self.ate_alpha = ate_alpha
         self.control_name = control_name
+
         self.propensity = None
         self.propensity_model = None
 
@@ -76,16 +96,26 @@ def __repr__(self):
         )
 
     def fit(self, X, treatment, y, p=None):
+        """Fit the inference model.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method; the
+                feature matrix is otherwise kept in its native format throughout.
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
+                the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+        """
+        X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
 
         if p is None:
-            # base.py does raw numpy indexing internally — convert at this boundary
-            self._set_propensity_models(
-                X=to_numpy(X), treatment=treatment_np, y=to_numpy(y)
-            )
+            self._set_propensity_models(X=X, treatment=treatment_np, y=to_numpy(y))
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -109,17 +139,15 @@ def fit(self, X, treatment, y, p=None):
             y_filt = filter_mask(y, mask)
             w = (to_numpy(treatment_filt) == group).astype(int)
 
-            self.models_mu_c[group].fit(
-                filter_mask(X_filt, w == 0), filter_mask(y_filt, w == 0)
-            )
-            self.models_mu_t[group].fit(
-                filter_mask(X_filt, w == 1), filter_mask(y_filt, w == 1)
-            )
-
-            y_filt_np = to_numpy(y_filt)
             X_filt_c = filter_mask(X_filt, w == 0)
             X_filt_t = filter_mask(X_filt, w == 1)
+            y_filt_np = to_numpy(y_filt)
 
+            # Train outcome models
+            self.models_mu_c[group].fit(X_filt_c, filter_mask(y_filt, w == 0))
+            self.models_mu_t[group].fit(X_filt_t, filter_mask(y_filt, w == 1))
+
+            # Calculate variances and treatment effects
             var_c = (
                 y_filt_np[w == 0] - self.models_mu_c[group].predict(X_filt_c)
             ).var()
@@ -129,6 +157,7 @@ def fit(self, X, treatment, y, p=None):
             ).var()
             self.vars_t[group] = var_t
 
+            # Train treatment models
             d_c = self.models_mu_t[group].predict(X_filt_c) - y_filt_np[w == 0]
             d_t = y_filt_np[w == 1] - self.models_mu_c[group].predict(X_filt_t)
             self.models_tau_c[group].fit(X_filt_c, d_c)
@@ -137,6 +166,23 @@ def fit(self, X, treatment, y, p=None):
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
+            y (np.array, pd.Series, or pl.Series, optional): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
+                the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X = collect_if_lazy(X)
+
         if p is None:
             logger.info("Generating propensity score")
             p = {
@@ -146,8 +192,7 @@ def predict(
         else:
             p = self._format_p(p, self.t_groups)
 
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         dhat_cs = {}
         dhat_ts = {}
 
@@ -196,6 +241,26 @@ def fit_predict(
         return_components=False,
         verbose=True,
     ):
+        """Fit the treatment effect and outcome models of the X learner and predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
+                the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            return_ci (bool): whether to return confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (str): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects. Output dim: [n_samples, n_treatment]
+                If return_ci, returns CATE [n_samples, n_treatment], LB [n_samples, n_treatment],
+                UB [n_samples, n_treatment]
+        """
+        X = collect_if_lazy(X)
         self.fit(X, treatment, y, p)
 
         if p is None:
@@ -210,7 +275,6 @@ def fit_predict(
         if not return_ci:
             return te
         else:
-            X_np = to_numpy(X)
             treatment_np = to_numpy(treatment)
             y_np = to_numpy(y)
 
@@ -221,12 +285,12 @@ def fit_predict(
             models_tau_c_global = deepcopy(self.models_tau_c)
             models_tau_t_global = deepcopy(self.models_tau_t)
             te_bootstraps = np.zeros(
-                shape=(X_np.shape[0], self.t_groups.shape[0], n_bootstraps)
+                shape=(n_rows(X), self.t_groups.shape[0], n_bootstraps)
             )
 
             logger.info("Bootstrap Confidence Intervals")
             for i in tqdm(range(n_bootstraps)):
-                te_b = self.bootstrap(X_np, treatment_np, y_np, p, size=bootstrap_size)
+                te_b = self.bootstrap(X, treatment_np, y_np, p, size=bootstrap_size)
                 te_bootstraps[:, :, i] = te_b
 
             te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
@@ -254,6 +318,24 @@ def estimate_ate(
         bootstrap_size=10000,
         pretrain=False,
     ):
+        """Estimate the Average Treatment Effect (ATE).
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
+                the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            bootstrap_ci (bool): whether run bootstrap for confidence intervals
+            n_bootstraps (int): number of bootstrap iterations
+            bootstrap_size (int): number of samples per bootstrap
+            pretrain (bool): whether a model has been fit, default False.
+        Returns:
+            The mean and confidence interval (LB, UB) of the ATE estimate.
+        """
+        X = collect_if_lazy(X)
+
         if pretrain:
             if p is None:
                 if not self.propensity:
@@ -294,6 +376,8 @@ def estimate_ate(
             dhat_t = dhat_ts[group][mask]
             p_filt = p[group][mask]
 
+            # SE formula is based on the lower bound formula (7) from Imbens, Guido W., and Jeffrey M. Wooldridge. 2009.
+            # "Recent Developments in the Econometrics of Program Evaluation." Journal of Economic Literature
             se = np.sqrt(
                 (
                     self.vars_t[group] / prob_treatment
@@ -313,7 +397,7 @@ def estimate_ate(
         if not bootstrap_ci:
             return ate, ate_lb, ate_ub
         else:
-            X_np = to_numpy(X)
+            y_np = to_numpy(y)
             t_groups_global = self.t_groups
             _classes_global = self._classes
             models_mu_c_global = deepcopy(self.models_mu_c)
@@ -325,9 +409,7 @@ def estimate_ate(
             ate_bootstraps = np.zeros(shape=(self.t_groups.shape[0], n_bootstraps))
 
             for n in tqdm(range(n_bootstraps)):
-                cate_b = self.bootstrap(
-                    X_np, treatment_np, to_numpy(y), p, size=bootstrap_size
-                )
+                cate_b = self.bootstrap(X, treatment_np, y_np, p, size=bootstrap_size)
                 ate_bootstraps[:, n] = cate_b.mean(axis=0)
 
             ate_lower = np.percentile(
@@ -347,6 +429,8 @@ def estimate_ate(
 
 
 class BaseXRegressor(BaseXLearner):
+    """A parent class for X-learner regressor classes."""
+
     def __init__(
         self,
         learner=None,
@@ -357,6 +441,18 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
+        """Initialize an X-learner regressor.
+
+        Args:
+            learner (optional): a model to estimate outcomes and treatment effects in both the control and treatment
+                groups
+            control_outcome_learner (optional): a model to estimate outcomes in the control group
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group
+            control_effect_learner (optional): a model to estimate treatment effects in the control group
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         super().__init__(
             learner=learner,
             control_outcome_learner=control_outcome_learner,
@@ -369,6 +465,8 @@ def __init__(
 
 
 class BaseXClassifier(BaseXLearner):
+    """A parent class for X-learner classifier classes."""
+
     def __init__(
         self,
         outcome_learner=None,
@@ -380,6 +478,24 @@ def __init__(
         ate_alpha=0.05,
         control_name=0,
     ):
+        """Initialize an X-learner classifier.
+
+        Args:
+            outcome_learner (optional): a model to estimate outcomes in both the control and treatment groups.
+                Should be a classifier.
+            effect_learner (optional): a model to estimate treatment effects in both the control and treatment groups.
+                Should be a regressor.
+            control_outcome_learner (optional): a model to estimate outcomes in the control group.
+                Should be a classifier.
+            treatment_outcome_learner (optional): a model to estimate outcomes in the treatment group.
+                Should be a classifier.
+            control_effect_learner (optional): a model to estimate treatment effects in the control group.
+                Should be a regressor.
+            treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
+                Should be a regressor.
+            ate_alpha (float, optional): the confidence level alpha of the ATE estimate
+            control_name (str or int, optional): name of control group
+        """
         if outcome_learner is not None:
             control_outcome_learner = outcome_learner
             treatment_outcome_learner = outcome_learner
@@ -405,16 +521,26 @@ def __init__(
             )
 
     def fit(self, X, treatment, y, p=None):
+        """Fit the inference model.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method; the
+                feature matrix is otherwise kept in its native format throughout.
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
+                the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+        """
+        X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
 
         if p is None:
-            # base.py does raw numpy indexing internally — convert at this boundary
-            self._set_propensity_models(
-                X=to_numpy(X), treatment=treatment_np, y=to_numpy(y)
-            )
+            self._set_propensity_models(X=X, treatment=treatment_np, y=to_numpy(y))
             p = self.propensity
         else:
             p = self._format_p(p, self.t_groups)
@@ -438,17 +564,15 @@ def fit(self, X, treatment, y, p=None):
             y_filt = filter_mask(y, mask)
             w = (to_numpy(treatment_filt) == group).astype(int)
 
-            self.models_mu_c[group].fit(
-                filter_mask(X_filt, w == 0), filter_mask(y_filt, w == 0)
-            )
-            self.models_mu_t[group].fit(
-                filter_mask(X_filt, w == 1), filter_mask(y_filt, w == 1)
-            )
-
-            y_filt_np = to_numpy(y_filt)
             X_filt_c = filter_mask(X_filt, w == 0)
             X_filt_t = filter_mask(X_filt, w == 1)
+            y_filt_np = to_numpy(y_filt)
+
+            # Train outcome models
+            self.models_mu_c[group].fit(X_filt_c, filter_mask(y_filt, w == 0))
+            self.models_mu_t[group].fit(X_filt_t, filter_mask(y_filt, w == 1))
 
+            # Calculate variances and treatment effects
             var_c = (
                 y_filt_np[w == 0]
                 - self.models_mu_c[group].predict_proba(X_filt_c)[:, 1]
@@ -460,6 +584,7 @@ def fit(self, X, treatment, y, p=None):
             ).var()
             self.vars_t[group] = var_t
 
+            # Train treatment models
             d_c = (
                 self.models_mu_t[group].predict_proba(X_filt_c)[:, 1]
                 - y_filt_np[w == 0]
@@ -474,6 +599,23 @@ def fit(self, X, treatment, y, p=None):
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
+        """Predict treatment effects.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
+                A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
+            y (np.array, pd.Series, or pl.Series, optional): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
+                the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
+            return_components (bool, optional): whether to return outcome for treatment and control seperately
+            verbose (bool, optional): whether to output progress logs
+        Returns:
+            (numpy.ndarray): Predictions of treatment effects.
+        """
+        X = collect_if_lazy(X)
+
         if p is None:
             logger.info("Generating propensity score")
             p = {
@@ -483,8 +625,7 @@ def predict(
         else:
             p = self._format_p(p, self.t_groups)
 
-        X_np = to_numpy(X)
-        te = np.zeros((X_np.shape[0], self.t_groups.shape[0]))
+        te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         dhat_cs = {}
         dhat_ts = {}
 
diff --git a/causalml/propensity.py b/causalml/propensity.py
index ec6aa981..42586b27 100644
--- a/causalml/propensity.py
+++ b/causalml/propensity.py
@@ -7,8 +7,6 @@
 from sklearn.isotonic import IsotonicRegression
 import xgboost as xgb
 
-from causalml.inference.meta.utils import convert_pd_to_np
-
 logger = logging.getLogger("causalml")
 
 
@@ -40,10 +38,11 @@ def fit(self, X, y):
         Fit a propensity model.
 
         Args:
-            X (numpy.ndarray): a feature matrix
-            y (numpy.ndarray): a binary target vector
+            X (numpy.ndarray, pd.DataFrame, or pl.DataFrame): a feature matrix.
+                scikit-learn >= 1.6 accepts pandas and Polars DataFrames
+                natively, so no conversion is performed here.
+            y (numpy.ndarray, pd.Series, or pl.Series): a binary target vector
         """
-        X, y = convert_pd_to_np(X, y)
         self.model.fit(X, y)
         if self.calibrate:
             # Fit a calibrator to the propensity scores with IsotonicRegression.
@@ -60,12 +59,11 @@ def predict(self, X):
         Predict propensity scores.
 
         Args:
-            X (numpy.ndarray): a feature matrix
+            X (numpy.ndarray, pd.DataFrame, or pl.DataFrame): a feature matrix
 
         Returns:
             (numpy.ndarray): Propensity scores between 0 and 1.
         """
-        X = convert_pd_to_np(X)
         p = self.model.predict_proba(X)[:, 1]
         if self.calibrate:
             p = self.calibrator.transform(p)
@@ -77,8 +75,8 @@ def fit_predict(self, X, y):
         Fit a propensity model and predict propensity scores.
 
         Args:
-            X (numpy.ndarray): a feature matrix
-            y (numpy.ndarray): a binary target vector
+            X (numpy.ndarray, pd.DataFrame, or pl.DataFrame): a feature matrix
+            y (numpy.ndarray, pd.Series, or pl.Series): a binary target vector
 
         Returns:
             (numpy.ndarray): Propensity scores between 0 and 1.
@@ -163,11 +161,9 @@ def fit(self, X, y, stop_val_size=0.2):
         Fit a propensity model.
 
         Args:
-            X (numpy.ndarray): a feature matrix
-            y (numpy.ndarray): a binary target vector
+            X (numpy.ndarray, pd.DataFrame, or pl.DataFrame): a feature matrix
+            y (numpy.ndarray, pd.Series, or pl.Series): a binary target vector
         """
-        X, y = convert_pd_to_np(X, y)
-
         if self.early_stop:
             X_train, X_val, y_train, y_val = train_test_split(
                 X, y, test_size=stop_val_size
@@ -201,29 +197,22 @@ def compute_propensity_score(
     """Generate propensity score if user didn't provide and optionally calibrate.
 
     Args:
-        X (np.matrix): features for training
-        treatment (np.array or pd.Series or pl.Series): a treatment vector for training
+        X (np.matrix, pd.DataFrame, or pl.DataFrame): features for training
+        treatment (np.array, pd.Series, or pl.Series): a treatment vector for training
         p_model (model object, optional): a binary classifier with either a predict_proba or predict method
-        X_pred (np.matrix, optional): features for prediction
-        treatment_pred (np.array or pd.Series or pl.Series, optional): a treatment vector for prediction
+        X_pred (np.matrix, pd.DataFrame, or pl.DataFrame, optional): features for prediction
+        treatment_pred (np.array, pd.Series, or pl.Series, optional): a treatment vector for prediction
         calibrate_p (bool, optional): whether calibrate the propensity score
-        clip_bounds (tuple, optional): lower and upper bounds for clipping propensity scores.
+        clip_bounds (tuple, optional): lower and upper bounds for clipping propensity scores. Bounds should be implemented
+                    such that: 0 < lower < upper < 1, to avoid division by zero in BaseRLearner.fit_predict() step.
 
     Returns:
         (tuple)
             - p (numpy.ndarray): propensity score
             - p_model (PropensityModel): either the original p_model or a trained ElasticNetPropensityModel
     """
-    # Normalise inputs to numpy so downstream sklearn models always see arrays.
-    X, treatment = convert_pd_to_np(X, treatment)
-
     if treatment_pred is None:
-        treatment_pred = treatment.copy()
-    else:
-        treatment_pred = convert_pd_to_np(treatment_pred)
-
-    if X_pred is not None:
-        X_pred = convert_pd_to_np(X_pred)
+        treatment_pred = treatment.copy() if hasattr(treatment, "copy") else treatment
 
     if p_model is None:
         p_model = ElasticNetPropensityModel(
@@ -240,4 +229,4 @@ def compute_propensity_score(
         logger.info("predict_proba not available, using predict instead")
         p = p_model.predict(X_pred)
 
-    return p, p_model
+    return p, p_model
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 883c42bc..9f6fa222 100755
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -44,9 +44,11 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+polars = ["polars>=1.0.0"]
 test = [
     "pytest>=4.6",
-    "pytest-cov>=4.0"
+    "pytest-cov>=4.0",
+    "causalml[polars]"
 ]
 tf = [
     "tensorflow>=2.4.0"

From d56b47e9d94f82635ca04924f98ec94147f318a0 Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Thu, 18 Jun 2026 23:06:32 +0530
Subject: [PATCH 05/15] fix lint

---
 causalml/inference/meta/tlearner.py | 175 ++++------------------------
 causalml/inference/meta/xlearner.py | 127 ++------------------
 causalml/propensity.py              |   2 +-
 3 files changed, 31 insertions(+), 273 deletions(-)

diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py
index 9efaa117..857cfa11 100644
--- a/causalml/inference/meta/tlearner.py
+++ b/causalml/inference/meta/tlearner.py
@@ -1,4 +1,3 @@
-import copy
 from copy import deepcopy
 import logging
 import numpy as np
@@ -62,35 +61,19 @@ def __init__(
         else:
             self.model_c = control_learner
 
-        # Preserve the unfitted template so repeated fit() calls always start fresh.
-        self._model_c_template = self.model_c
-
         if treatment_learner is None:
             self.model_t = deepcopy(learner)
         else:
             self.model_t = treatment_learner
 
-        # Preserve the unfitted template so repeated fit() calls always start fresh.
-        self._model_t_template = self.model_t
-
         self.ate_alpha = ate_alpha
         self.control_name = control_name
-        self.bootstrap_models_ = None
 
     def __repr__(self):
         return "{}(model_c={}, model_t={})".format(
             self.__class__.__name__, self.model_c.__repr__(), self.model_t.__repr__()
         )
 
-    def _unfitted_clone(self):
-        template = copy.copy(self)
-        for attr in ("models_c", "models_t", "bootstrap_models_"):
-            if hasattr(template, attr):
-                delattr(template, attr)
-        template.model_c = self._model_c_template
-        template.model_t = self._model_t_template
-        return template
-
     @ignore_warnings(category=ConvergenceWarning)
     def fit(self, X, treatment, y, p=None):
         """Fit the inference model.
@@ -101,36 +84,6 @@ def fit(self, X, treatment, y, p=None):
                 feature matrix is otherwise kept in its native format throughout.
             treatment (np.array, pd.Series, or pl.Series): a treatment vector
             y (np.array, pd.Series, or pl.Series): an outcome vector
-    def fit(
-        self,
-        X,
-        treatment,
-        y,
-        p=None,
-        store_bootstraps=False,
-        n_bootstraps=200,
-        bootstrap_size=10000,
-        random_state=None,
-        n_jobs=1,
-    ):
-        """Fit the inference model
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p: unused, kept for API consistency
-            store_bootstraps (bool, optional): if True, trains a bootstrap ensemble
-                during fit and stores it in self.bootstrap_models_ for post-fit CI
-                estimation via predict(return_ci=True). Default: False.
-            n_bootstraps (int, optional): number of bootstrap iterations. Default: 200.
-                Note: storing N bootstraps of a GBM-based learner with k treatment
-                groups holds 2*N*k model objects in memory. Monitor RAM for large N
-                or heavy base learners.
-            n_jobs (int, optional): number of parallel jobs for bootstrap fitting.
-                -1 uses all available cores. Default: 1.
-            bootstrap_size (int, optional): number of samples per bootstrap. Default: 10000.
-            random_state (int, optional): random seed for reproducible bootstrap sampling.
         """
         X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
@@ -138,17 +91,9 @@ def fit(
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
+        self.models_c = {group: deepcopy(self.model_c) for group in self.t_groups}
         self.models_t = {group: deepcopy(self.model_t) for group in self.t_groups}
 
-        # model_c is trained on the control group, which is identical for every
-        # treatment group, so fit it once. Deepcopy from the unfitted template so
-        # re-calling fit() always starts from a clean state (safe with warm_start).
-        control_mask = treatment == self.control_name
-        self.model_c = deepcopy(self._model_c_template)
-        self.model_c.fit(X[control_mask], y[control_mask])
-        # Expose as a shared-reference dict to preserve the public models_c API.
-        self.models_c = {group: self.model_c for group in self.t_groups}
-
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
             treatment_filt = filter_mask(treatment, mask)
@@ -162,52 +107,9 @@ def fit(
             self.models_t[group].fit(
                 filter_mask(X_filt, w == 1), filter_mask(y_filt, w == 1)
             )
-            treatment_mask = treatment == group
-            self.models_t[group].fit(X[treatment_mask], y[treatment_mask])
-
-        if store_bootstraps:
-            self.fit_bootstrap_ensemble(
-                X=X,
-                treatment=treatment,
-                y=y,
-                n_bootstraps=n_bootstraps,
-                bootstrap_size=bootstrap_size,
-                random_state=random_state,
-                n_jobs=n_jobs,
-            )
-        else:
-            self.bootstrap_models_ = None
-
-    def _compute_bootstrap_ci(self, X):
-        """Compute bootstrap CI using stored ensemble.
-
-        Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-        Returns:
-            (te_lower, te_upper): percentile CI bounds, each of shape [n_samples, n_treatment]
-        """
-        if self.bootstrap_models_ is None:
-            raise ValueError(
-                "No bootstrap ensemble found. Call fit(..., store_bootstraps=True) first."
-            )
-        te_bootstraps = np.zeros(
-            (X.shape[0], self.t_groups.shape[0], len(self.bootstrap_models_))
-        )
-        for b, learner_b in enumerate(self.bootstrap_models_):
-            te_bootstraps[:, :, b] = learner_b.predict(X)
-        te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
-        te_upper = np.percentile(te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2)
-        return te_lower, te_upper
 
     def predict(
-        self,
-        X,
-        treatment=None,
-        y=None,
-        p=None,
-        return_components=False,
-        verbose=True,
-        return_ci=False,
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
         """Predict treatment effects.
 
@@ -217,36 +119,19 @@ def predict(
             treatment (np.array, pd.Series, or pl.Series, optional): a treatment vector
             y (np.array, pd.Series, or pl.Series, optional): an outcome vector
             return_components (bool, optional): whether to return outcome for treatment and control seperately
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
-            return_components (bool, optional): whether to return outcome for
-                treatment and control separately
             verbose (bool, optional): whether to output progress logs
-            return_ci (bool, optional): whether to return confidence intervals
-                using the stored bootstrap ensemble. Requires fit() to have been
-                called with store_bootstraps=True. CI width is controlled by
-                self.ate_alpha set at init time.
         Returns:
-            (numpy.ndarray): Predictions of treatment effects. If return_ci=True,
-                returns (te, te_lower, te_upper) each of shape [n_samples, n_treatment].
-                return_ci=True and return_components=True cannot be used together.
+            (numpy.ndarray): Predictions of treatment effects.
         """
         X = collect_if_lazy(X)
         yhat_cs = {}
-        if return_ci and return_components:
-            raise ValueError("return_ci and return_components cannot both be True.")
-
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         yhat_ts = {}
 
-        yhat_c = self.model_c.predict(X)
-        # Build a shared-reference dict so return_components callers keep the
-        # yhat_cs[group] indexing API without duplicating the underlying array.
-        yhat_cs = {group: yhat_c for group in self.t_groups}
-
         for group in self.t_groups:
-            yhat_ts[group] = self.models_t[group].predict(X)
+            model_c = self.models_c[group]
+            model_t = self.models_t[group]
+            yhat_cs[group] = model_c.predict(X)
+            yhat_ts[group] = model_t.predict(X)
 
             if (y is not None) and (treatment is not None) and verbose:
                 treatment_np = to_numpy(treatment)
@@ -256,7 +141,7 @@ def predict(
                 w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
-                yhat[w == 0] = yhat_c[mask][w == 0]
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
                 yhat[w == 1] = yhat_ts[group][mask][w == 1]
 
                 logger.info("Error metrics for group {}".format(group))
@@ -264,11 +149,7 @@ def predict(
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
-            te[:, i] = yhat_ts[group] - yhat_c
-
-        if return_ci:
-            te_lower, te_upper = self._compute_bootstrap_ci(X)
-            return te, te_lower, te_upper
+            te[:, i] = yhat_ts[group] - yhat_cs[group]
 
         if not return_components:
             return te
@@ -315,7 +196,7 @@ def fit_predict(
 
             t_groups_global = self.t_groups
             _classes_global = self._classes
-            model_c_global = deepcopy(self.model_c)
+            models_c_global = deepcopy(self.models_c)
             models_t_global = deepcopy(self.models_t)
             te_bootstraps = np.zeros(
                 shape=(n_rows(X), self.t_groups.shape[0], n_bootstraps)
@@ -334,8 +215,7 @@ def fit_predict(
             # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
-            self.model_c = deepcopy(model_c_global)
-            self.models_c = {group: self.model_c for group in self.t_groups}
+            self.models_c = deepcopy(models_c_global)
             self.models_t = deepcopy(models_t_global)
 
             return (te, te_lower, te_upper)
@@ -412,7 +292,7 @@ def estimate_ate(
         else:
             t_groups_global = self.t_groups
             _classes_global = self._classes
-            model_c_global = deepcopy(self.model_c)
+            models_c_global = deepcopy(self.models_c)
             models_t_global = deepcopy(self.models_t)
 
             logger.info("Bootstrap Confidence Intervals for ATE")
@@ -432,8 +312,7 @@ def estimate_ate(
             # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
-            self.model_c = deepcopy(model_c_global)
-            self.models_c = {group: self.model_c for group in self.t_groups}
+            self.models_c = deepcopy(models_c_global)
             self.models_t = deepcopy(models_t_global)
 
             return ate, ate_lower, ate_upper
@@ -497,14 +376,7 @@ def __init__(
         )
 
     def predict(
-        self,
-        X,
-        treatment=None,
-        y=None,
-        p=None,
-        return_components=False,
-        verbose=True,
-        return_ci=False,
+        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
         """Predict treatment effects.
 
@@ -520,16 +392,13 @@ def predict(
         """
         X = collect_if_lazy(X)
         yhat_cs = {}
-        if return_ci and return_components:
-            raise ValueError("return_ci and return_components cannot both be True.")
-
         yhat_ts = {}
 
-        yhat_c = self.model_c.predict_proba(X)[:, 1]
-        yhat_cs = {group: yhat_c for group in self.t_groups}
-
         for group in self.t_groups:
-            yhat_ts[group] = self.models_t[group].predict_proba(X)[:, 1]
+            model_c = self.models_c[group]
+            model_t = self.models_t[group]
+            yhat_cs[group] = model_c.predict_proba(X)[:, 1]
+            yhat_ts[group] = model_t.predict_proba(X)[:, 1]
 
             if (y is not None) and (treatment is not None) and verbose:
                 treatment_np = to_numpy(treatment)
@@ -539,7 +408,7 @@ def predict(
                 w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
-                yhat[w == 0] = yhat_c[mask][w == 0]
+                yhat[w == 0] = yhat_cs[group][mask][w == 0]
                 yhat[w == 1] = yhat_ts[group][mask][w == 1]
 
                 logger.info("Error metrics for group {}".format(group))
@@ -547,11 +416,7 @@ def predict(
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
-            te[:, i] = yhat_ts[group] - yhat_c
-
-        if return_ci:
-            te_lower, te_upper = self._compute_bootstrap_ci(X)
-            return te, te_lower, te_upper
+            te[:, i] = yhat_ts[group] - yhat_cs[group]
 
         if not return_components:
             return te
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index e77da8fb..f12e886d 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -74,28 +74,6 @@ def __init__(
             if treatment_effect_learner is None
             else treatment_effect_learner
         )
-        if control_outcome_learner is None:
-            self.model_mu_c = deepcopy(learner)
-        else:
-            self.model_mu_c = control_outcome_learner
-
-        # Preserve the unfitted template so repeated fit() calls always start fresh.
-        self._model_mu_c_template = self.model_mu_c
-
-        if treatment_outcome_learner is None:
-            self.model_mu_t = deepcopy(learner)
-        else:
-            self.model_mu_t = treatment_outcome_learner
-
-        if control_effect_learner is None:
-            self.model_tau_c = deepcopy(learner)
-        else:
-            self.model_tau_c = control_effect_learner
-
-        if treatment_effect_learner is None:
-            self.model_tau_t = deepcopy(learner)
-        else:
-            self.model_tau_t = treatment_effect_learner
 
         self.ate_alpha = ate_alpha
         self.control_name = control_name
@@ -143,6 +121,7 @@ def fit(self, X, treatment, y, p=None):
             p = self._format_p(p, self.t_groups)
 
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
+        self.models_mu_c = {group: deepcopy(self.model_mu_c) for group in self.t_groups}
         self.models_mu_t = {group: deepcopy(self.model_mu_t) for group in self.t_groups}
         self.models_tau_c = {
             group: deepcopy(self.model_tau_c) for group in self.t_groups
@@ -153,20 +132,6 @@ def fit(self, X, treatment, y, p=None):
         self.vars_c = {}
         self.vars_t = {}
 
-        # model_mu_c is trained on control data, which is the same for every treatment
-        # group. Deepcopy from the unfitted template so re-calling fit() starts fresh.
-        control_mask = treatment == self.control_name
-        self.model_mu_c = deepcopy(self._model_mu_c_template)
-        self.model_mu_c.fit(X[control_mask], y[control_mask])
-        # Expose as a shared-reference dict to preserve the public models_mu_c API.
-        self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
-
-        # var_c depends only on model_mu_c and control data — constant across groups.
-        y_control_pred = self.model_mu_c.predict(X[control_mask])
-        self.var_c = (y[control_mask] - y_control_pred).var()
-        # Keep vars_c dict for backward compatibility with existing callers.
-        self.vars_c = {group: self.var_c for group in self.t_groups}
-
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
             treatment_filt = filter_mask(treatment, mask)
@@ -189,26 +154,14 @@ def fit(self, X, treatment, y, p=None):
             self.vars_c[group] = var_c
             var_t = (
                 y_filt_np[w == 1] - self.models_mu_t[group].predict(X_filt_t)
-            treatment_mask = treatment == group
-            X_treat = X[treatment_mask]
-            y_treat = y[treatment_mask]
-
-            self.models_mu_t[group].fit(X_treat, y_treat)
-
-            self.vars_t[group] = (
-                y_treat - self.models_mu_t[group].predict(X_treat)
             ).var()
+            self.vars_t[group] = var_t
 
             # Train treatment models
             d_c = self.models_mu_t[group].predict(X_filt_c) - y_filt_np[w == 0]
             d_t = y_filt_np[w == 1] - self.models_mu_c[group].predict(X_filt_t)
             self.models_tau_c[group].fit(X_filt_c, d_c)
             self.models_tau_t[group].fit(X_filt_t, d_t)
-            # Train treatment effect models using cross-group imputation
-            d_c = self.models_mu_t[group].predict(X[control_mask]) - y[control_mask]
-            d_t = y_treat - self.model_mu_c.predict(X_treat)
-            self.models_tau_c[group].fit(X[control_mask], d_c)
-            self.models_tau_t[group].fit(X_treat, d_t)
 
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
@@ -243,12 +196,6 @@ def predict(
         dhat_cs = {}
         dhat_ts = {}
 
-        # For verbose metrics, control predictions are constant across groups.
-        yhat_c_verbose = None
-        if (y is not None) and (treatment is not None) and verbose:
-            control_mask = treatment == self.control_name
-            yhat_c_verbose = self.model_mu_c.predict(X[control_mask])
-
         for i, group in enumerate(self.t_groups):
             dhat_cs[group] = self.models_tau_c[group].predict(X)
             dhat_ts[group] = self.models_tau_t[group].predict(X)
@@ -273,16 +220,6 @@ def predict(
                 yhat[w == 1] = self.models_mu_t[group].predict(
                     filter_mask(X_filt, w == 1)
                 )
-            if yhat_c_verbose is not None:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                X_filt = X[mask]
-                y_filt = y[mask]
-                w = (treatment_filt == group).astype(int)
-
-                yhat = np.zeros_like(y_filt, dtype=float)
-                yhat[w == 0] = yhat_c_verbose
-                yhat[w == 1] = self.models_mu_t[group].predict(X_filt[w == 1])
 
                 logger.info("Error metrics for group {}".format(group))
                 regression_metrics(y_filt, yhat, w)
@@ -343,7 +280,7 @@ def fit_predict(
 
             t_groups_global = self.t_groups
             _classes_global = self._classes
-            model_mu_c_global = deepcopy(self.model_mu_c)
+            models_mu_c_global = deepcopy(self.models_mu_c)
             models_mu_t_global = deepcopy(self.models_mu_t)
             models_tau_c_global = deepcopy(self.models_tau_c)
             models_tau_t_global = deepcopy(self.models_tau_t)
@@ -363,8 +300,7 @@ def fit_predict(
 
             self.t_groups = t_groups_global
             self._classes = _classes_global
-            self.model_mu_c = deepcopy(model_mu_c_global)
-            self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
+            self.models_mu_c = deepcopy(models_mu_c_global)
             self.models_mu_t = deepcopy(models_mu_t_global)
             self.models_tau_c = deepcopy(models_tau_c_global)
             self.models_tau_t = deepcopy(models_tau_t_global)
@@ -445,7 +381,7 @@ def estimate_ate(
             se = np.sqrt(
                 (
                     self.vars_t[group] / prob_treatment
-                    + self.var_c / (1 - prob_treatment)
+                    + self.vars_c[group] / (1 - prob_treatment)
                     + (p_filt * dhat_c + (1 - p_filt) * dhat_t).var()
                 )
                 / w.shape[0]
@@ -464,7 +400,7 @@ def estimate_ate(
             y_np = to_numpy(y)
             t_groups_global = self.t_groups
             _classes_global = self._classes
-            model_mu_c_global = deepcopy(self.model_mu_c)
+            models_mu_c_global = deepcopy(self.models_mu_c)
             models_mu_t_global = deepcopy(self.models_mu_t)
             models_tau_c_global = deepcopy(self.models_tau_c)
             models_tau_t_global = deepcopy(self.models_tau_t)
@@ -485,8 +421,7 @@ def estimate_ate(
 
             self.t_groups = t_groups_global
             self._classes = _classes_global
-            self.model_mu_c = deepcopy(model_mu_c_global)
-            self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
+            self.models_mu_c = deepcopy(models_mu_c_global)
             self.models_mu_t = deepcopy(models_mu_t_global)
             self.models_tau_c = deepcopy(models_tau_c_global)
             self.models_tau_t = deepcopy(models_tau_t_global)
@@ -611,6 +546,7 @@ def fit(self, X, treatment, y, p=None):
             p = self._format_p(p, self.t_groups)
 
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
+        self.models_mu_c = {group: deepcopy(self.model_mu_c) for group in self.t_groups}
         self.models_mu_t = {group: deepcopy(self.model_mu_t) for group in self.t_groups}
         self.models_tau_c = {
             group: deepcopy(self.model_tau_c) for group in self.t_groups
@@ -621,18 +557,6 @@ def fit(self, X, treatment, y, p=None):
         self.vars_c = {}
         self.vars_t = {}
 
-        # model_mu_c is trained on control data, which is the same for every treatment
-        # group, so fit it once and store as a single model (not a per-group dict).
-        control_mask = treatment == self.control_name
-        self.model_mu_c = deepcopy(self._model_mu_c_template)
-        self.model_mu_c.fit(X[control_mask], y[control_mask])
-        self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
-
-        # var_c depends only on model_mu_c and control data — constant across groups.
-        y_control_pred = self.model_mu_c.predict_proba(X[control_mask])[:, 1]
-        self.var_c = (y[control_mask] - y_control_pred).var()
-        self.vars_c = {group: self.var_c for group in self.t_groups}
-
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
             treatment_filt = filter_mask(treatment, mask)
@@ -657,17 +581,10 @@ def fit(self, X, treatment, y, p=None):
             var_t = (
                 y_filt_np[w == 1]
                 - self.models_mu_t[group].predict_proba(X_filt_t)[:, 1]
-            treatment_mask = treatment == group
-            X_treat = X[treatment_mask]
-            y_treat = y[treatment_mask]
-
-            self.models_mu_t[group].fit(X_treat, y_treat)
-
-            self.vars_t[group] = (
-                y_treat - self.models_mu_t[group].predict_proba(X_treat)[:, 1]
             ).var()
+            self.vars_t[group] = var_t
 
-            # Train treatment effect models using cross-group imputation
+            # Train treatment models
             d_c = (
                 self.models_mu_t[group].predict_proba(X_filt_c)[:, 1]
                 - y_filt_np[w == 0]
@@ -678,12 +595,6 @@ def fit(self, X, treatment, y, p=None):
             )
             self.models_tau_c[group].fit(X_filt_c, d_c)
             self.models_tau_t[group].fit(X_filt_t, d_t)
-                self.models_mu_t[group].predict_proba(X[control_mask])[:, 1]
-                - y[control_mask]
-            )
-            d_t = y_treat - self.model_mu_c.predict_proba(X_treat)[:, 1]
-            self.models_tau_c[group].fit(X[control_mask], d_c)
-            self.models_tau_t[group].fit(X_treat, d_t)
 
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
@@ -718,12 +629,6 @@ def predict(
         dhat_cs = {}
         dhat_ts = {}
 
-        # For verbose metrics, control predictions are constant across groups.
-        yhat_c_verbose = None
-        if (y is not None) and (treatment is not None) and verbose:
-            control_mask = treatment == self.control_name
-            yhat_c_verbose = self.model_mu_c.predict_proba(X[control_mask])[:, 1]
-
         for i, group in enumerate(self.t_groups):
             dhat_cs[group] = self.models_tau_c[group].predict(X)
             dhat_ts[group] = self.models_tau_t[group].predict(X)
@@ -748,18 +653,6 @@ def predict(
                 yhat[w == 1] = self.models_mu_t[group].predict_proba(
                     filter_mask(X_filt, w == 1)
                 )[:, 1]
-            if yhat_c_verbose is not None:
-                mask = (treatment == group) | (treatment == self.control_name)
-                treatment_filt = treatment[mask]
-                X_filt = X[mask]
-                y_filt = y[mask]
-                w = (treatment_filt == group).astype(int)
-
-                yhat = np.zeros_like(y_filt, dtype=float)
-                yhat[w == 0] = yhat_c_verbose
-                yhat[w == 1] = self.models_mu_t[group].predict_proba(X_filt[w == 1])[
-                    :, 1
-                ]
 
                 logger.info("Error metrics for group {}".format(group))
                 classification_metrics(y_filt, yhat, w)
diff --git a/causalml/propensity.py b/causalml/propensity.py
index 42586b27..fc0e0640 100644
--- a/causalml/propensity.py
+++ b/causalml/propensity.py
@@ -229,4 +229,4 @@ def compute_propensity_score(
         logger.info("predict_proba not available, using predict instead")
         p = p_model.predict(X_pred)
 
-    return p, p_model
\ No newline at end of file
+    return p, p_model

From c87523c775cdfb8370f2f9b47adfee180c59521a Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Sat, 20 Jun 2026 12:08:49 +0530
Subject: [PATCH 06/15] fixing the build errors

---
 causalml/inference/meta/tlearner.py | 196 ++++++++++++++++++++++------
 causalml/inference/meta/utils.py    |   6 +-
 2 files changed, 158 insertions(+), 44 deletions(-)

diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py
index 857cfa11..3467ad9e 100644
--- a/causalml/inference/meta/tlearner.py
+++ b/causalml/inference/meta/tlearner.py
@@ -1,3 +1,4 @@
+import copy
 from copy import deepcopy
 import logging
 import numpy as np
@@ -61,21 +62,48 @@ def __init__(
         else:
             self.model_c = control_learner
 
+        # Preserve the unfitted template so repeated fit() calls always start fresh.
+        self._model_c_template = self.model_c
+
         if treatment_learner is None:
             self.model_t = deepcopy(learner)
         else:
             self.model_t = treatment_learner
 
+        # Preserve the unfitted template so repeated fit() calls always start fresh.
+        self._model_t_template = self.model_t
+
         self.ate_alpha = ate_alpha
         self.control_name = control_name
+        self.bootstrap_models_ = None
 
     def __repr__(self):
         return "{}(model_c={}, model_t={})".format(
             self.__class__.__name__, self.model_c.__repr__(), self.model_t.__repr__()
         )
 
+    def _unfitted_clone(self):
+        template = copy.copy(self)
+        for attr in ("models_c", "models_t", "bootstrap_models_"):
+            if hasattr(template, attr):
+                delattr(template, attr)
+        template.model_c = self._model_c_template
+        template.model_t = self._model_t_template
+        return template
+
     @ignore_warnings(category=ConvergenceWarning)
-    def fit(self, X, treatment, y, p=None):
+    def fit(
+        self,
+        X,
+        treatment,
+        y,
+        p=None,
+        store_bootstraps=False,
+        n_bootstraps=200,
+        bootstrap_size=10000,
+        random_state=None,
+        n_jobs=1,
+    ):
         """Fit the inference model.
 
         Args:
@@ -84,32 +112,84 @@ def fit(self, X, treatment, y, p=None):
                 feature matrix is otherwise kept in its native format throughout.
             treatment (np.array, pd.Series, or pl.Series): a treatment vector
             y (np.array, pd.Series, or pl.Series): an outcome vector
+            p: unused, kept for API consistency
+            store_bootstraps (bool, optional): if True, trains a bootstrap ensemble
+                during fit and stores it in self.bootstrap_models_ for post-fit CI
+                estimation via predict(return_ci=True). Default: False.
+            n_bootstraps (int, optional): number of bootstrap iterations. Default: 200.
+            n_jobs (int, optional): number of parallel jobs for bootstrap fitting.
+                -1 uses all available cores. Default: 1.
+            bootstrap_size (int, optional): number of samples per bootstrap. Default: 10000.
+            random_state (int, optional): random seed for reproducible bootstrap sampling.
         """
         X = collect_if_lazy(X)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
         self.t_groups.sort()
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
-        self.models_c = {group: deepcopy(self.model_c) for group in self.t_groups}
         self.models_t = {group: deepcopy(self.model_t) for group in self.t_groups}
 
+        # model_c is trained on the control group, which is identical for every
+        # treatment group, so fit it once. Deepcopy from the unfitted template so
+        # re-calling fit() always starts from a clean state.
+        control_mask = treatment_np == self.control_name
+        self.model_c = deepcopy(self._model_c_template)
+        self.model_c.fit(filter_mask(X, control_mask), y_np[control_mask])
+        # Expose as a shared-reference dict to preserve the public models_c API.
+        self.models_c = {group: self.model_c for group in self.t_groups}
+
         for group in self.t_groups:
-            mask = (treatment_np == group) | (treatment_np == self.control_name)
-            treatment_filt = filter_mask(treatment, mask)
-            X_filt = filter_mask(X, mask)
-            y_filt = filter_mask(y, mask)
-            w = (to_numpy(treatment_filt) == group).astype(int)
+            treatment_mask = treatment_np == group
+            self.models_t[group].fit(
+                filter_mask(X, treatment_mask), y_np[treatment_mask]
+            )
 
-            self.models_c[group].fit(
-                filter_mask(X_filt, w == 0), filter_mask(y_filt, w == 0)
+        if store_bootstraps:
+            self.fit_bootstrap_ensemble(
+                X=X,
+                treatment=treatment_np,
+                y=y_np,
+                n_bootstraps=n_bootstraps,
+                bootstrap_size=bootstrap_size,
+                random_state=random_state,
+                n_jobs=n_jobs,
             )
-            self.models_t[group].fit(
-                filter_mask(X_filt, w == 1), filter_mask(y_filt, w == 1)
+        else:
+            self.bootstrap_models_ = None
+
+    def _compute_bootstrap_ci(self, X):
+        """Compute bootstrap CI using stored ensemble.
+
+        Args:
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+        Returns:
+            (te_lower, te_upper): percentile CI bounds, each of shape [n_samples, n_treatment]
+        """
+        if self.bootstrap_models_ is None:
+            raise ValueError(
+                "No bootstrap ensemble found. Call fit(..., store_bootstraps=True) first."
             )
+        te_bootstraps = np.zeros(
+            (n_rows(X), self.t_groups.shape[0], len(self.bootstrap_models_))
+        )
+        for b, learner_b in enumerate(self.bootstrap_models_):
+            te_bootstraps[:, :, b] = learner_b.predict(X)
+        te_lower = np.percentile(te_bootstraps, (self.ate_alpha / 2) * 100, axis=2)
+        te_upper = np.percentile(te_bootstraps, (1 - self.ate_alpha / 2) * 100, axis=2)
+        return te_lower, te_upper
 
     def predict(
-        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+        self,
+        X,
+        treatment=None,
+        y=None,
+        p=None,
+        return_components=False,
+        verbose=True,
+        return_ci=False,
     ):
         """Predict treatment effects.
 
@@ -120,18 +200,25 @@ def predict(
             y (np.array, pd.Series, or pl.Series, optional): an outcome vector
             return_components (bool, optional): whether to return outcome for treatment and control seperately
             verbose (bool, optional): whether to output progress logs
+            return_ci (bool, optional): whether to return confidence intervals using
+                the stored bootstrap ensemble. Requires fit() to have been called
+                with store_bootstraps=True.
         Returns:
-            (numpy.ndarray): Predictions of treatment effects.
+            (numpy.ndarray): Predictions of treatment effects. If return_ci=True,
+                returns (te, te_lower, te_upper) each of shape [n_samples, n_treatment].
         """
+        if return_ci and return_components:
+            raise ValueError("return_ci and return_components cannot both be True.")
+
         X = collect_if_lazy(X)
-        yhat_cs = {}
         yhat_ts = {}
 
+        yhat_c = self.model_c.predict(X)
+        # Shared-reference dict — no array duplication
+        yhat_cs = {group: yhat_c for group in self.t_groups}
+
         for group in self.t_groups:
-            model_c = self.models_c[group]
-            model_t = self.models_t[group]
-            yhat_cs[group] = model_c.predict(X)
-            yhat_ts[group] = model_t.predict(X)
+            yhat_ts[group] = self.models_t[group].predict(X)
 
             if (y is not None) and (treatment is not None) and verbose:
                 treatment_np = to_numpy(treatment)
@@ -141,7 +228,7 @@ def predict(
                 w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
-                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 0] = yhat_c[mask][w == 0]
                 yhat[w == 1] = yhat_ts[group][mask][w == 1]
 
                 logger.info("Error metrics for group {}".format(group))
@@ -149,7 +236,11 @@ def predict(
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
-            te[:, i] = yhat_ts[group] - yhat_cs[group]
+            te[:, i] = yhat_ts[group] - yhat_c
+
+        if return_ci:
+            te_lower, te_upper = self._compute_bootstrap_ci(X)
+            return te, te_lower, te_upper
 
         if not return_components:
             return te
@@ -185,18 +276,18 @@ def fit_predict(
                 UB [n_samples, n_treatment]
         """
         X = collect_if_lazy(X)
-        self.fit(X, treatment, y)
-        te = self.predict(X, treatment, y, return_components=return_components)
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
+        self.fit(X, treatment_np, y_np)
+        te = self.predict(X, treatment_np, y_np, return_components=return_components)
 
         if not return_ci:
             return te
         else:
-            treatment_np = to_numpy(treatment)
-            y_np = to_numpy(y)
-
             t_groups_global = self.t_groups
             _classes_global = self._classes
-            models_c_global = deepcopy(self.models_c)
+            model_c_global = deepcopy(self.model_c)
             models_t_global = deepcopy(self.models_t)
             te_bootstraps = np.zeros(
                 shape=(n_rows(X), self.t_groups.shape[0], n_bootstraps)
@@ -215,7 +306,8 @@ def fit_predict(
             # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
-            self.models_c = deepcopy(models_c_global)
+            self.model_c = deepcopy(model_c_global)
+            self.models_c = {group: self.model_c for group in self.t_groups}
             self.models_t = deepcopy(models_t_global)
 
             return (te, te_lower, te_upper)
@@ -245,16 +337,18 @@ def estimate_ate(
             The mean and confidence interval (LB, UB) of the ATE estimate.
         """
         X = collect_if_lazy(X)
+        treatment_np = to_numpy(treatment)
+        y_np = to_numpy(y)
+
         if pretrain:
-            te, yhat_cs, yhat_ts = self.predict(X, treatment, y, return_components=True)
+            te, yhat_cs, yhat_ts = self.predict(
+                X, treatment_np, y_np, return_components=True
+            )
         else:
             te, yhat_cs, yhat_ts = self.fit_predict(
-                X, treatment, y, return_components=True
+                X, treatment_np, y_np, return_components=True
             )
 
-        treatment_np = to_numpy(treatment)
-        y_np = to_numpy(y)
-
         ate = np.zeros(self.t_groups.shape[0])
         ate_lb = np.zeros(self.t_groups.shape[0])
         ate_ub = np.zeros(self.t_groups.shape[0])
@@ -292,7 +386,7 @@ def estimate_ate(
         else:
             t_groups_global = self.t_groups
             _classes_global = self._classes
-            models_c_global = deepcopy(self.models_c)
+            model_c_global = deepcopy(self.model_c)
             models_t_global = deepcopy(self.models_t)
 
             logger.info("Bootstrap Confidence Intervals for ATE")
@@ -312,7 +406,8 @@ def estimate_ate(
             # set member variables back to global (currently last bootstrapped outcome)
             self.t_groups = t_groups_global
             self._classes = _classes_global
-            self.models_c = deepcopy(models_c_global)
+            self.model_c = deepcopy(model_c_global)
+            self.models_c = {group: self.model_c for group in self.t_groups}
             self.models_t = deepcopy(models_t_global)
 
             return ate, ate_lower, ate_upper
@@ -376,7 +471,14 @@ def __init__(
         )
 
     def predict(
-        self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
+        self,
+        X,
+        treatment=None,
+        y=None,
+        p=None,
+        return_components=False,
+        verbose=True,
+        return_ci=False,
     ):
         """Predict treatment effects.
 
@@ -387,18 +489,22 @@ def predict(
             y (np.array, pd.Series, or pl.Series, optional): an outcome vector
             return_components (bool, optional): whether to return outcome for treatment and control seperately
             verbose (bool, optional): whether to output progress logs
+            return_ci (bool, optional): whether to return confidence intervals using
+                the stored bootstrap ensemble.
         Returns:
             (numpy.ndarray): Predictions of treatment effects.
         """
+        if return_ci and return_components:
+            raise ValueError("return_ci and return_components cannot both be True.")
+
         X = collect_if_lazy(X)
-        yhat_cs = {}
         yhat_ts = {}
 
+        yhat_c = self.model_c.predict_proba(X)[:, 1]
+        yhat_cs = {group: yhat_c for group in self.t_groups}
+
         for group in self.t_groups:
-            model_c = self.models_c[group]
-            model_t = self.models_t[group]
-            yhat_cs[group] = model_c.predict_proba(X)[:, 1]
-            yhat_ts[group] = model_t.predict_proba(X)[:, 1]
+            yhat_ts[group] = self.models_t[group].predict_proba(X)[:, 1]
 
             if (y is not None) and (treatment is not None) and verbose:
                 treatment_np = to_numpy(treatment)
@@ -408,15 +514,21 @@ def predict(
                 w = (treatment_filt_np == group).astype(int)
 
                 yhat = np.zeros_like(y_filt, dtype=float)
-                yhat[w == 0] = yhat_cs[group][mask][w == 0]
+                yhat[w == 0] = yhat_c[mask][w == 0]
                 yhat[w == 1] = yhat_ts[group][mask][w == 1]
 
                 logger.info("Error metrics for group {}".format(group))
+                from causalml.metrics import classification_metrics
+
                 classification_metrics(y_filt, yhat, w)
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
-            te[:, i] = yhat_ts[group] - yhat_cs[group]
+            te[:, i] = yhat_ts[group] - yhat_c
+
+        if return_ci:
+            te_lower, te_upper = self._compute_bootstrap_ci(X)
+            return te, te_lower, te_upper
 
         if not return_components:
             return te
diff --git a/causalml/inference/meta/utils.py b/causalml/inference/meta/utils.py
index 14b39b7c..88b8dc0a 100644
--- a/causalml/inference/meta/utils.py
+++ b/causalml/inference/meta/utils.py
@@ -152,7 +152,8 @@ def prepend_column(value, X):
     n = n_rows(X)
 
     if isinstance(X, pd.DataFrame):
-        col = pd.DataFrame({"_w": np.full(n, value)}, index=X.index)
+        col_name = X.columns[0].__class__(0) if len(X.columns) > 0 else "_w"
+        col = pd.DataFrame({col_name: np.full(n, value)}, index=X.index)
         return pd.concat([col, X], axis=1)
 
     if _POLARS_AVAILABLE and isinstance(X, pl.DataFrame):
@@ -181,7 +182,8 @@ def concat_treatment_col(w, X):
     X = collect_if_lazy(X)
 
     if isinstance(X, pd.DataFrame):
-        col = pd.DataFrame({"_w": np.asarray(w)}, index=X.index)
+        col_name = X.columns[0].__class__(0) if len(X.columns) > 0 else "_w"
+        col = pd.DataFrame({col_name: np.asarray(w)}, index=X.index)
         return pd.concat([col, X], axis=1)
 
     if _POLARS_AVAILABLE and isinstance(X, pl.DataFrame):

From 863df3ab14be1734ea11cf282ff876325edca0d8 Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Sat, 20 Jun 2026 12:33:25 +0530
Subject: [PATCH 07/15] shared-ref models_mu_c in XLearner and pandas column
 name clash in S-learner

---
 causalml/inference/meta/utils.py    | 13 +++++++------
 causalml/inference/meta/xlearner.py | 18 ++++++++++--------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/causalml/inference/meta/utils.py b/causalml/inference/meta/utils.py
index 88b8dc0a..91dff021 100644
--- a/causalml/inference/meta/utils.py
+++ b/causalml/inference/meta/utils.py
@@ -152,9 +152,11 @@ def prepend_column(value, X):
     n = n_rows(X)
 
     if isinstance(X, pd.DataFrame):
-        col_name = X.columns[0].__class__(0) if len(X.columns) > 0 else "_w"
-        col = pd.DataFrame({col_name: np.full(n, value)}, index=X.index)
-        return pd.concat([col, X], axis=1)
+        # Reset to range-index column names to avoid type clashes, then restore.
+        # We prepend a treatment column then rename all columns to 0..n so
+        # downstream sklearn sees consistent integer names.
+        arr = np.hstack((np.full((n, 1), value), X.to_numpy()))
+        return pd.DataFrame(arr, index=X.index)
 
     if _POLARS_AVAILABLE and isinstance(X, pl.DataFrame):
         col = pl.Series("_w", np.full(n, value))
@@ -182,9 +184,8 @@ def concat_treatment_col(w, X):
     X = collect_if_lazy(X)
 
     if isinstance(X, pd.DataFrame):
-        col_name = X.columns[0].__class__(0) if len(X.columns) > 0 else "_w"
-        col = pd.DataFrame({col_name: np.asarray(w)}, index=X.index)
-        return pd.concat([col, X], axis=1)
+        arr = np.hstack((np.asarray(w).reshape(-1, 1), X.to_numpy()))
+        return pd.DataFrame(arr, index=X.index)
 
     if _POLARS_AVAILABLE and isinstance(X, pl.DataFrame):
         col = pl.Series("_w", np.asarray(w))
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index f12e886d..ddc08b20 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -121,7 +121,6 @@ def fit(self, X, treatment, y, p=None):
             p = self._format_p(p, self.t_groups)
 
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
-        self.models_mu_c = {group: deepcopy(self.model_mu_c) for group in self.t_groups}
         self.models_mu_t = {group: deepcopy(self.model_mu_t) for group in self.t_groups}
         self.models_tau_c = {
             group: deepcopy(self.model_tau_c) for group in self.t_groups
@@ -132,6 +131,12 @@ def fit(self, X, treatment, y, p=None):
         self.vars_c = {}
         self.vars_t = {}
 
+        # model_mu_c is trained on control only (identical across groups) - fit once.
+        control_mask = treatment_np == self.control_name
+        self.model_mu_c = deepcopy(self.model_mu_c)
+        self.model_mu_c.fit(filter_mask(X, control_mask), filter_mask(y, control_mask))
+        self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
+
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
             treatment_filt = filter_mask(treatment, mask)
@@ -143,23 +148,20 @@ def fit(self, X, treatment, y, p=None):
             X_filt_t = filter_mask(X_filt, w == 1)
             y_filt_np = to_numpy(y_filt)
 
-            # Train outcome models
-            self.models_mu_c[group].fit(X_filt_c, filter_mask(y_filt, w == 0))
+            # Train treatment outcome model
             self.models_mu_t[group].fit(X_filt_t, filter_mask(y_filt, w == 1))
 
             # Calculate variances and treatment effects
-            var_c = (
-                y_filt_np[w == 0] - self.models_mu_c[group].predict(X_filt_c)
-            ).var()
+            var_c = (y_filt_np[w == 0] - self.model_mu_c.predict(X_filt_c)).var()
             self.vars_c[group] = var_c
             var_t = (
                 y_filt_np[w == 1] - self.models_mu_t[group].predict(X_filt_t)
             ).var()
             self.vars_t[group] = var_t
 
-            # Train treatment models
+            # Train treatment effect models
             d_c = self.models_mu_t[group].predict(X_filt_c) - y_filt_np[w == 0]
-            d_t = y_filt_np[w == 1] - self.models_mu_c[group].predict(X_filt_t)
+            d_t = y_filt_np[w == 1] - self.model_mu_c.predict(X_filt_t)
             self.models_tau_c[group].fit(X_filt_c, d_c)
             self.models_tau_t[group].fit(X_filt_t, d_t)
 

From 7e90d5def77d507c361f8fff910f32285f149b94 Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Sat, 20 Jun 2026 20:17:00 +0530
Subject: [PATCH 08/15] expose var_c as finite scalar on BaseXLearner

---
 causalml/inference/meta/xlearner.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index ddc08b20..a59a25fe 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -128,14 +128,19 @@ def fit(self, X, treatment, y, p=None):
         self.models_tau_t = {
             group: deepcopy(self.model_tau_t) for group in self.t_groups
         }
-        self.vars_c = {}
         self.vars_t = {}
 
-        # model_mu_c is trained on control only (identical across groups) - fit once.
+        # model_mu_c is trained on control only (identical across groups) — fit once.
         control_mask = treatment_np == self.control_name
+        X_control = filter_mask(X, control_mask)
+        y_control = to_numpy(filter_mask(y, control_mask))
         self.model_mu_c = deepcopy(self.model_mu_c)
-        self.model_mu_c.fit(filter_mask(X, control_mask), filter_mask(y, control_mask))
+        self.model_mu_c.fit(X_control, y_control)
         self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
+        # var_c is a single scalar since control model is shared across groups
+        self.var_c = (y_control - self.model_mu_c.predict(X_control)).var()
+        # Keep vars_c dict for backward compat with estimate_ate
+        self.vars_c = {group: self.var_c for group in self.t_groups}
 
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
@@ -151,9 +156,6 @@ def fit(self, X, treatment, y, p=None):
             # Train treatment outcome model
             self.models_mu_t[group].fit(X_filt_t, filter_mask(y_filt, w == 1))
 
-            # Calculate variances and treatment effects
-            var_c = (y_filt_np[w == 0] - self.model_mu_c.predict(X_filt_c)).var()
-            self.vars_c[group] = var_c
             var_t = (
                 y_filt_np[w == 1] - self.models_mu_t[group].predict(X_filt_t)
             ).var()

From 04d8194e1aba57df1446aa2a6b73c2d4f616b41f Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Sun, 21 Jun 2026 13:40:39 +0530
Subject: [PATCH 09/15] fixing the review changes

---
 causalml/inference/meta/drlearner.py |   4 +-
 causalml/inference/meta/slearner.py  |   4 --
 causalml/inference/meta/tlearner.py  |   1 -
 causalml/inference/meta/xlearner.py  |  23 +++---
 tests/test_polars_support.py         | 100 +++++++++++++++++++++++++++
 5 files changed, 115 insertions(+), 17 deletions(-)

diff --git a/causalml/inference/meta/drlearner.py b/causalml/inference/meta/drlearner.py
index 828c19cf..4855cfe1 100644
--- a/causalml/inference/meta/drlearner.py
+++ b/causalml/inference/meta/drlearner.py
@@ -271,7 +271,7 @@ def predict(
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         yhat_cs = {}
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+
         yhat_ts = {}
 
         # models_mu_c is fold-specific but not group-specific; predict once and reuse.
@@ -633,7 +633,7 @@ def predict(
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
         yhat_cs = {}
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
+
         yhat_ts = {}
 
         # models_mu_c is fold-specific but not group-specific; predict once and reuse.
diff --git a/causalml/inference/meta/slearner.py b/causalml/inference/meta/slearner.py
index 52ed0b38..0e93dc60 100644
--- a/causalml/inference/meta/slearner.py
+++ b/causalml/inference/meta/slearner.py
@@ -129,8 +129,6 @@ def predict(
         # Build the augmented arrays once; they are identical for every group.
         # (Separate allocations avoid in-place mutation by learners like CatBoost
         # that set the writeable flag to False on arrays passed to predict().)
-        X_new_c = np.hstack((np.zeros((X.shape[0], 1)), X))
-        X_new_t = np.hstack((np.ones((X.shape[0], 1)), X))
 
         for group in self.t_groups:
             model = self.models[group]
@@ -139,7 +137,6 @@ def predict(
             # mutation, which fails when learners like CatBoost set the
             # writeable flag to False on arrays passed to predict().
             X_new_c = prepend_column(0.0, X)
-            yhat_cs[group] = model.predict(X_new_c)
 
             X_new_t = prepend_column(1.0, X)
             yhat_cs[group] = model.predict(X_new_c)
@@ -389,7 +386,6 @@ def predict(
             model = self.models[group]
 
             X_new_c = prepend_column(0.0, X)
-            yhat_cs[group] = model.predict_proba(X_new_c)[:, 1]
 
             X_new_t = prepend_column(1.0, X)
             yhat_cs[group] = model.predict_proba(X_new_c)[:, 1]
diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py
index 3467ad9e..57266fab 100644
--- a/causalml/inference/meta/tlearner.py
+++ b/causalml/inference/meta/tlearner.py
@@ -518,7 +518,6 @@ def predict(
                 yhat[w == 1] = yhat_ts[group][mask][w == 1]
 
                 logger.info("Error metrics for group {}".format(group))
-                from causalml.metrics import classification_metrics
 
                 classification_metrics(y_filt, yhat, w)
 
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index a59a25fe..b95c4f3a 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -550,7 +550,6 @@ def fit(self, X, treatment, y, p=None):
             p = self._format_p(p, self.t_groups)
 
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
-        self.models_mu_c = {group: deepcopy(self.model_mu_c) for group in self.t_groups}
         self.models_mu_t = {group: deepcopy(self.model_mu_t) for group in self.t_groups}
         self.models_tau_c = {
             group: deepcopy(self.model_tau_c) for group in self.t_groups
@@ -558,9 +557,20 @@ def fit(self, X, treatment, y, p=None):
         self.models_tau_t = {
             group: deepcopy(self.model_tau_t) for group in self.t_groups
         }
-        self.vars_c = {}
         self.vars_t = {}
 
+        # model_mu_c is trained on control only (identical across groups) — fit once.
+        control_mask = treatment_np == self.control_name
+        X_control = filter_mask(X, control_mask)
+        y_control = filter_mask(y, control_mask)
+        self.model_mu_c = deepcopy(self.model_mu_c)
+        self.model_mu_c.fit(X_control, y_control)
+        self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
+        self.var_c = (
+            to_numpy(y_control) - self.model_mu_c.predict_proba(X_control)[:, 1]
+        ).var()
+        self.vars_c = {group: self.var_c for group in self.t_groups}
+
         for group in self.t_groups:
             mask = (treatment_np == group) | (treatment_np == self.control_name)
             treatment_filt = filter_mask(treatment, mask)
@@ -572,16 +582,9 @@ def fit(self, X, treatment, y, p=None):
             X_filt_t = filter_mask(X_filt, w == 1)
             y_filt_np = to_numpy(y_filt)
 
-            # Train outcome models
-            self.models_mu_c[group].fit(X_filt_c, filter_mask(y_filt, w == 0))
+            # Train treatment outcome model
             self.models_mu_t[group].fit(X_filt_t, filter_mask(y_filt, w == 1))
 
-            # Calculate variances and treatment effects
-            var_c = (
-                y_filt_np[w == 0]
-                - self.models_mu_c[group].predict_proba(X_filt_c)[:, 1]
-            ).var()
-            self.vars_c[group] = var_c
             var_t = (
                 y_filt_np[w == 1]
                 - self.models_mu_t[group].predict_proba(X_filt_t)[:, 1]
diff --git a/tests/test_polars_support.py b/tests/test_polars_support.py
index 949cd91f..2e50de4a 100644
--- a/tests/test_polars_support.py
+++ b/tests/test_polars_support.py
@@ -15,6 +15,10 @@
 
 from sklearn.linear_model import LinearRegression
 
+from sklearn.linear_model import LogisticRegression
+from causalml.inference.meta.tlearner import BaseTClassifier
+from causalml.inference.meta.slearner import BaseSClassifier
+from causalml.inference.meta.xlearner import BaseXClassifier
 from causalml.inference.meta.tlearner import BaseTRegressor
 from causalml.inference.meta.slearner import BaseSRegressor
 from causalml.inference.meta.xlearner import BaseXRegressor
@@ -324,3 +328,99 @@ def test_polars_predict_only(self, synthetic_data_numpy, synthetic_data_polars):
         te = learner.predict(X_pl)
         assert isinstance(te, np.ndarray)
         assert te.shape[0] == N
+
+
+# T-Learner Classifier
+
+
+class TestTClassifierPolars:
+    @pytest.fixture(autouse=True)
+    def _learner(self):
+        self.learner = BaseTClassifier(learner=LogisticRegression())
+
+    def _fit_predict(self, X, treatment, y):
+        self.learner.fit(X, treatment, y)
+        return self.learner.predict(X)
+
+    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseTClassifier(learner=LogisticRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_np, te_pl)
+
+    def test_lazyframe_input(self, synthetic_data_numpy, synthetic_data_polars_lazy):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseTClassifier(learner=LogisticRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars_lazy)
+        _assert_te_close(te_np, te_pl)
+
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
+        te = self._fit_predict(*synthetic_data_polars)
+        assert isinstance(te, np.ndarray)
+
+
+# S-Learner Classifier
+
+
+class TestSClassifierPolars:
+    @pytest.fixture(autouse=True)
+    def _learner(self):
+        self.learner = BaseSClassifier(learner=LogisticRegression())
+
+    def _fit_predict(self, X, treatment, y):
+        self.learner.fit(X, treatment, y)
+        return self.learner.predict(X)
+
+    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseSClassifier(learner=LogisticRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_np, te_pl)
+
+    def test_lazyframe_input(self, synthetic_data_numpy, synthetic_data_polars_lazy):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseSClassifier(learner=LogisticRegression())
+        te_pl = self._fit_predict(*synthetic_data_polars_lazy)
+        _assert_te_close(te_np, te_pl)
+
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
+        te = self._fit_predict(*synthetic_data_polars)
+        assert isinstance(te, np.ndarray)
+
+
+# X-Learner Classifier
+
+
+class TestXClassifierPolars:
+    @pytest.fixture(autouse=True)
+    def _learner(self):
+        self.learner = BaseXClassifier(
+            outcome_learner=LogisticRegression(),
+            effect_learner=LinearRegression(),
+        )
+
+    def _fit_predict(self, X, treatment, y):
+        self.learner.fit(X, treatment, y)
+        return self.learner.predict(X)
+
+    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseXClassifier(
+            outcome_learner=LogisticRegression(),
+            effect_learner=LinearRegression(),
+        )
+        te_pl = self._fit_predict(*synthetic_data_polars)
+        _assert_te_close(te_np, te_pl)
+
+    def test_lazyframe_input(self, synthetic_data_numpy, synthetic_data_polars_lazy):
+        te_np = self._fit_predict(*synthetic_data_numpy)
+        self.learner = BaseXClassifier(
+            outcome_learner=LogisticRegression(),
+            effect_learner=LinearRegression(),
+        )
+        te_pl = self._fit_predict(*synthetic_data_polars_lazy)
+        _assert_te_close(te_np, te_pl)
+
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
+        te = self._fit_predict(*synthetic_data_polars)
+        assert isinstance(te, np.ndarray)

From 511bd5fa5c5b0a0ea7ca903441e1b2dc7fb653c5 Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Sun, 21 Jun 2026 13:50:04 +0530
Subject: [PATCH 10/15] update the tests

---
 tests/test_polars_support.py | 85 ++++++++++++++++++++++++++----------
 1 file changed, 61 insertions(+), 24 deletions(-)

diff --git a/tests/test_polars_support.py b/tests/test_polars_support.py
index 2e50de4a..8b2c1f70 100644
--- a/tests/test_polars_support.py
+++ b/tests/test_polars_support.py
@@ -70,6 +70,31 @@ def synthetic_data_polars_lazy(synthetic_data_polars):
     return X.lazy(), treatment, y
 
 
+@pytest.fixture(scope="module")
+def synthetic_data_numpy_binary():
+    """Return (X, treatment, y) with binary y for classifier tests."""
+    rng = np.random.default_rng(RANDOM_STATE)
+    X = rng.standard_normal((N, N_FEATURES))
+    treatment = rng.choice([0, 1], size=N)
+    y = rng.choice([0, 1], size=N)
+    return X, treatment, y
+
+
+@pytest.fixture(scope="module")
+def synthetic_data_polars_binary(synthetic_data_numpy_binary):
+    X_np, t_np, y_np = synthetic_data_numpy_binary
+    X = pl.DataFrame({f"f{i}": X_np[:, i] for i in range(N_FEATURES)})
+    treatment = pl.Series("treatment", t_np)
+    y = pl.Series("outcome", y_np)
+    return X, treatment, y
+
+
+@pytest.fixture(scope="module")
+def synthetic_data_polars_lazy_binary(synthetic_data_polars_binary):
+    X, treatment, y = synthetic_data_polars_binary
+    return X.lazy(), treatment, y
+
+
 # convert_pd_to_np unit tests
 
 
@@ -342,20 +367,24 @@ def _fit_predict(self, X, treatment, y):
         self.learner.fit(X, treatment, y)
         return self.learner.predict(X)
 
-    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
-        te_np = self._fit_predict(*synthetic_data_numpy)
+    def test_polars_matches_numpy(
+        self, synthetic_data_numpy_binary, synthetic_data_polars_binary
+    ):
+        te_np = self._fit_predict(*synthetic_data_numpy_binary)
         self.learner = BaseTClassifier(learner=LogisticRegression())
-        te_pl = self._fit_predict(*synthetic_data_polars)
+        te_pl = self._fit_predict(*synthetic_data_polars_binary)
         _assert_te_close(te_np, te_pl)
 
-    def test_lazyframe_input(self, synthetic_data_numpy, synthetic_data_polars_lazy):
-        te_np = self._fit_predict(*synthetic_data_numpy)
+    def test_lazyframe_input(
+        self, synthetic_data_numpy_binary, synthetic_data_polars_lazy_binary
+    ):
+        te_np = self._fit_predict(*synthetic_data_numpy_binary)
         self.learner = BaseTClassifier(learner=LogisticRegression())
-        te_pl = self._fit_predict(*synthetic_data_polars_lazy)
+        te_pl = self._fit_predict(*synthetic_data_polars_lazy_binary)
         _assert_te_close(te_np, te_pl)
 
-    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
-        te = self._fit_predict(*synthetic_data_polars)
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars_binary):
+        te = self._fit_predict(*synthetic_data_polars_binary)
         assert isinstance(te, np.ndarray)
 
 
@@ -371,20 +400,24 @@ def _fit_predict(self, X, treatment, y):
         self.learner.fit(X, treatment, y)
         return self.learner.predict(X)
 
-    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
-        te_np = self._fit_predict(*synthetic_data_numpy)
+    def test_polars_matches_numpy(
+        self, synthetic_data_numpy_binary, synthetic_data_polars_binary
+    ):
+        te_np = self._fit_predict(*synthetic_data_numpy_binary)
         self.learner = BaseSClassifier(learner=LogisticRegression())
-        te_pl = self._fit_predict(*synthetic_data_polars)
+        te_pl = self._fit_predict(*synthetic_data_polars_binary)
         _assert_te_close(te_np, te_pl)
 
-    def test_lazyframe_input(self, synthetic_data_numpy, synthetic_data_polars_lazy):
-        te_np = self._fit_predict(*synthetic_data_numpy)
+    def test_lazyframe_input(
+        self, synthetic_data_numpy_binary, synthetic_data_polars_lazy_binary
+    ):
+        te_np = self._fit_predict(*synthetic_data_numpy_binary)
         self.learner = BaseSClassifier(learner=LogisticRegression())
-        te_pl = self._fit_predict(*synthetic_data_polars_lazy)
+        te_pl = self._fit_predict(*synthetic_data_polars_lazy_binary)
         _assert_te_close(te_np, te_pl)
 
-    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
-        te = self._fit_predict(*synthetic_data_polars)
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars_binary):
+        te = self._fit_predict(*synthetic_data_polars_binary)
         assert isinstance(te, np.ndarray)
 
 
@@ -403,24 +436,28 @@ def _fit_predict(self, X, treatment, y):
         self.learner.fit(X, treatment, y)
         return self.learner.predict(X)
 
-    def test_polars_matches_numpy(self, synthetic_data_numpy, synthetic_data_polars):
-        te_np = self._fit_predict(*synthetic_data_numpy)
+    def test_polars_matches_numpy(
+        self, synthetic_data_numpy_binary, synthetic_data_polars_binary
+    ):
+        te_np = self._fit_predict(*synthetic_data_numpy_binary)
         self.learner = BaseXClassifier(
             outcome_learner=LogisticRegression(),
             effect_learner=LinearRegression(),
         )
-        te_pl = self._fit_predict(*synthetic_data_polars)
+        te_pl = self._fit_predict(*synthetic_data_polars_binary)
         _assert_te_close(te_np, te_pl)
 
-    def test_lazyframe_input(self, synthetic_data_numpy, synthetic_data_polars_lazy):
-        te_np = self._fit_predict(*synthetic_data_numpy)
+    def test_lazyframe_input(
+        self, synthetic_data_numpy_binary, synthetic_data_polars_lazy_binary
+    ):
+        te_np = self._fit_predict(*synthetic_data_numpy_binary)
         self.learner = BaseXClassifier(
             outcome_learner=LogisticRegression(),
             effect_learner=LinearRegression(),
         )
-        te_pl = self._fit_predict(*synthetic_data_polars_lazy)
+        te_pl = self._fit_predict(*synthetic_data_polars_lazy_binary)
         _assert_te_close(te_np, te_pl)
 
-    def test_fit_predict_returns_numpy(self, synthetic_data_polars):
-        te = self._fit_predict(*synthetic_data_polars)
+    def test_fit_predict_returns_numpy(self, synthetic_data_polars_binary):
+        te = self._fit_predict(*synthetic_data_polars_binary)
         assert isinstance(te, np.ndarray)

From 661e8a82a0cc5eacc132618f052e43876f7c9c8f Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Mon, 29 Jun 2026 12:20:50 +0530
Subject: [PATCH 11/15] fixing merge errors

---
 causalml/inference/meta/drlearner.py | 16 +++-------------
 causalml/inference/meta/rlearner.py  | 14 +++++++++-----
 causalml/inference/meta/xlearner.py  | 14 +++-----------
 3 files changed, 15 insertions(+), 29 deletions(-)

diff --git a/causalml/inference/meta/drlearner.py b/causalml/inference/meta/drlearner.py
index a0a9c0da..46712ea1 100644
--- a/causalml/inference/meta/drlearner.py
+++ b/causalml/inference/meta/drlearner.py
@@ -72,7 +72,7 @@ def __init__(
             if treatment_effect_learner is None
             else treatment_effect_learner
         )
-
+        """
         Note: arguments are stored verbatim (scikit-learn convention) so that
         ``get_params`` / ``clone`` work correctly. Model construction is deferred to ``fit()``.
         Per the scikit-learn convention, ``__init__`` does not validate or raise —
@@ -103,7 +103,7 @@ def fit(self, X, treatment, y, p=None, seed=None):
                 single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
                 float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
             seed (int): random seed for cross-fitting
-        """
+
         X = collect_if_lazy(X)
             X (np.matrix or np.array or pd.Dataframe): a feature matrix
             treatment (np.array or pd.Series): a treatment vector
@@ -351,17 +351,7 @@ def fit_predict(
         verbose=True,
         seed=None,
     ):
-        """Fit the treatment effect and outcome models of the DR learner and predict treatment effects.
-
-        Args:
-            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
-            treatment (np.array, pd.Series, or pl.Series): a treatment vector
-            y (np.array, pd.Series, or pl.Series): an outcome vector
-            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
-                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-        """Fit the DR-learner and predict treatment effects.
-
+        """
         Args:
             X (np.matrix or np.array or pd.Dataframe): a feature matrix
             treatment (np.array or pd.Series): a treatment vector
diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py
index f99a24c0..abd198f2 100644
--- a/causalml/inference/meta/rlearner.py
+++ b/causalml/inference/meta/rlearner.py
@@ -90,7 +90,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
                 weight of each observation for `effect_learner`. If None, it assumes equal weight.
             verbose (bool, optional): whether to output progress logs
-        """
+
         X = collect_if_lazy(X)
             X (np.matrix or np.array or pd.Dataframe): a feature matrix
             treatment (np.array or pd.Series): a treatment vector
@@ -120,7 +120,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             sample_weight = to_numpy(sample_weight)
 
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
-            sample_weight = convert_pd_to_np(sample_weight)
+        sample_weight = convert_pd_to_np(sample_weight)
 
         self.t_groups = np.unique(treatment[treatment != self.control_name])
         self.t_groups.sort()
@@ -338,9 +338,13 @@ def estimate_ate(
                 or y_np is None
                 or not len(y_np)
             ):
-            if not len(treatment) or not len(y):
-                raise ValueError("treatment and y must be provided when pretrain=False")
-            te = self.fit_predict(X, treatment, y, p, sample_weight, return_ci=False)
+                if not len(treatment) or not len(y):
+                    raise ValueError(
+                        "treatment and y must be provided when pretrain=False"
+                    )
+                te = self.fit_predict(
+                    X, treatment, y, p, sample_weight, return_ci=False
+                )
 
         ate = np.zeros(self.t_groups.shape[0])
         ate_lb = np.zeros(self.t_groups.shape[0])
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index 125e64d6..d3b45f6b 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -74,7 +74,7 @@ def __init__(
             if treatment_effect_learner is None
             else treatment_effect_learner
         )
-
+        """
         Note: arguments are stored verbatim (scikit-learn convention) so that
         ``get_params`` / ``clone`` work correctly. Model construction is deferred to ``fit()``.
         Per the scikit-learn convention, ``__init__`` does not validate or raise —
@@ -302,16 +302,8 @@ def fit_predict(
         return_components=False,
         verbose=True,
     ):
-        """Fit the treatment effect and outcome models of the X learner and predict treatment effects.
-
-        Args:
-            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
-            treatment (np.array, pd.Series, or pl.Series): a treatment vector
-            y (np.array, pd.Series, or pl.Series): an outcome vector
-            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
-                the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
-                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-        """Fit the X-learner and predict treatment effects.
+        """
+        Fit the X-learner and predict treatment effects.
 
         Args:
             X (np.matrix or np.array or pd.Dataframe): a feature matrix

From 126ef38616c624d70899b9eea81e1d79f42ddd37 Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Mon, 29 Jun 2026 13:53:00 +0530
Subject: [PATCH 12/15] fixing the related errors

---
 causalml/inference/meta/drlearner.py |  1 +
 causalml/inference/meta/rlearner.py  |  1 +
 causalml/inference/meta/tlearner.py  |  5 +-
 causalml/inference/meta/xlearner.py  | 68 ++++++++++++++--------------
 4 files changed, 39 insertions(+), 36 deletions(-)

diff --git a/causalml/inference/meta/drlearner.py b/causalml/inference/meta/drlearner.py
index 46712ea1..942a56f3 100644
--- a/causalml/inference/meta/drlearner.py
+++ b/causalml/inference/meta/drlearner.py
@@ -16,6 +16,7 @@
     filter_index,
     n_rows,
     to_numpy,
+    convert_pd_to_np,
 )
 from causalml.metrics import regression_metrics, classification_metrics
 from causalml.propensity import compute_propensity_score
diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py
index abd198f2..c89725ed 100644
--- a/causalml/inference/meta/rlearner.py
+++ b/causalml/inference/meta/rlearner.py
@@ -15,6 +15,7 @@
     to_numpy,
     get_xgboost_objective_metric,
     get_weighted_variance,
+    convert_pd_to_np,
 )
 from causalml.propensity import ElasticNetPropensityModel
 
diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py
index c11b7439..cd0487be 100644
--- a/causalml/inference/meta/tlearner.py
+++ b/causalml/inference/meta/tlearner.py
@@ -21,6 +21,7 @@
     filter_mask,
     n_rows,
     to_numpy,
+    convert_pd_to_np,
 )
 from causalml.metrics import regression_metrics, classification_metrics
 
@@ -115,9 +116,7 @@ def fit(
         # model_c is trained on the control group, which is identical for every
         # treatment group, so fit it once. Deepcopy from the unfitted template so
         # re-calling fit() always starts from a clean state.
-        control_mask = treatment_np == self.control_name
-        self.model_c = deepcopy(self._model_c_template)
-        self.model_c.fit(filter_mask(X, control_mask), y_np[control_mask])
+
         # Resolve base models from stored constructor args (no templates needed).
         _control_learner = (
             self.control_learner
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index d3b45f6b..b93cb2fe 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -11,6 +11,7 @@
     filter_mask,
     n_rows,
     to_numpy,
+    convert_pd_to_np,
 )
 from causalml.metrics import regression_metrics, classification_metrics
 
@@ -209,10 +210,6 @@ def fit(self, X, treatment, y, p=None):
             d_t = y_filt_np[w == 1] - self.model_mu_c.predict(X_filt_t)
             self.models_tau_c[group].fit(X_filt_c, d_c)
             self.models_tau_t[group].fit(X_filt_t, d_t)
-            d_c = self.models_mu_t[group].predict(X[control_mask]) - y[control_mask]
-            d_t = y_treat - self.model_mu_c.predict(X_treat)
-            self.models_tau_c[group].fit(X[control_mask], d_c)
-            self.models_tau_t[group].fit(X_treat, d_t)
         return self
 
     def predict(
@@ -550,23 +547,6 @@ def __init__(
         # Sentinel so estimate_ate(pretrain=True) raises cleanly before fit().
         self.propensity = {}
 
-    def fit(self, X, treatment, y, p=None):
-        """Fit the inference model (classifier variant — uses predict_proba)."""
-        # Resolve and validate here (not in __init__) — sklearn convention.
-        _control_outcome_learner = self.control_outcome_learner or self.outcome_learner
-        _treatment_outcome_learner = (
-            self.treatment_outcome_learner or self.outcome_learner
-        )
-        _control_effect_learner = self.control_effect_learner or self.effect_learner
-        _treatment_effect_learner = self.treatment_effect_learner or self.effect_learner
-
-        if (
-            _control_outcome_learner is None or _treatment_outcome_learner is None
-        ) and (_control_effect_learner is None or _treatment_effect_learner is None):
-            raise ValueError(
-                "Either the outcome learner or the effect learner pair must be specified."
-            )
-
     def fit(self, X, treatment, y, p=None):
         """Fit the inference model.
 
@@ -581,6 +561,39 @@ def fit(self, X, treatment, y, p=None):
                 float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
         """
         X = collect_if_lazy(X)
+        if (self.outcome_learner is None) and (
+            (self.control_outcome_learner is None)
+            or (self.treatment_outcome_learner is None)
+            or (self.control_effect_learner is None)
+            or (self.treatment_effect_learner is None)
+        ):
+            raise ValueError(
+                "Either `outcome_learner` and `effect_learner`, or all four specialized learners must be specified."
+            )
+
+        _control_outcome_learner = (
+            deepcopy(self.outcome_learner)
+            if self.control_outcome_learner is None
+            else deepcopy(self.control_outcome_learner)
+        )
+
+        _treatment_outcome_learner = (
+            deepcopy(self.outcome_learner)
+            if self.treatment_outcome_learner is None
+            else deepcopy(self.treatment_outcome_learner)
+        )
+
+        _control_effect_learner = (
+            deepcopy(self.effect_learner)
+            if self.control_effect_learner is None
+            else deepcopy(self.control_effect_learner)
+        )
+
+        _treatment_effect_learner = (
+            deepcopy(self.effect_learner)
+            if self.treatment_effect_learner is None
+            else deepcopy(self.treatment_effect_learner)
+        )
         X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
@@ -607,15 +620,6 @@ def fit(self, X, treatment, y, p=None):
         self.vars_t = {}
 
         # model_mu_c is trained on control only (identical across groups) — fit once.
-        control_mask = treatment_np == self.control_name
-        X_control = filter_mask(X, control_mask)
-        y_control = filter_mask(y, control_mask)
-        self.model_mu_c = deepcopy(self.model_mu_c)
-        self.model_mu_c.fit(X_control, y_control)
-        self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
-        self.var_c = (
-            to_numpy(y_control) - self.model_mu_c.predict_proba(X_control)[:, 1]
-        ).var()
         control_mask = treatment == self.control_name
         self.model_mu_c = deepcopy(_control_outcome_learner)
         self.model_mu_c.fit(X[control_mask], y[control_mask])
@@ -656,9 +660,7 @@ def fit(self, X, treatment, y, p=None):
             )
             self.models_tau_c[group].fit(X_filt_c, d_c)
             self.models_tau_t[group].fit(X_filt_t, d_t)
-            d_t = y_treat - self.model_mu_c.predict_proba(X_treat)[:, 1]
-            self.models_tau_c[group].fit(X[control_mask], d_c)
-            self.models_tau_t[group].fit(X_treat, d_t)
+
         return self
 
     def predict(

From 8a935370568f2762946f8f4cfc8f32bab6686a6f Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Mon, 29 Jun 2026 14:27:45 +0530
Subject: [PATCH 13/15] checking CI errors

---
 causalml/inference/meta/rlearner.py | 39 ++++++++++-------------------
 1 file changed, 13 insertions(+), 26 deletions(-)

diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py
index c89725ed..f0aff405 100644
--- a/causalml/inference/meta/rlearner.py
+++ b/causalml/inference/meta/rlearner.py
@@ -92,7 +92,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
                 weight of each observation for `effect_learner`. If None, it assumes equal weight.
             verbose (bool, optional): whether to output progress logs
 
-        X = collect_if_lazy(X)
             X (np.matrix or np.array or pd.Dataframe): a feature matrix
             treatment (np.array or pd.Series): a treatment vector
             y (np.array or pd.Series): an outcome vector
@@ -100,6 +99,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             sample_weight (np.array or pd.Series, optional): sample weights for `effect_learner`.
             verbose (bool, optional): whether to output progress logs
         """
+        X = collect_if_lazy(X)
         if (self.learner is None) and (
             (self.outcome_learner is None) or (self.effect_learner is None)
         ):
@@ -121,9 +121,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             sample_weight = to_numpy(sample_weight)
 
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
-        sample_weight = convert_pd_to_np(sample_weight)
-
-        self.t_groups = np.unique(treatment[treatment != self.control_name])
         self.t_groups.sort()
 
         if p is None:
@@ -175,12 +172,7 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             diff_t = y_filt[w == 1] - yhat_filt[w == 1]
             if sample_weight is not None:
                 sample_weight_filt = sample_weight[mask]
-                self.vars_c[group] = get_weighted_variance(
-                    diff_c, sample_weight_filt[w == 0]
-                )
-                self.vars_t[group] = get_weighted_variance(
-                    diff_t, sample_weight_filt[w == 1]
-                )
+
                 sample_weight_filt_c = sample_weight_filt[w == 0]
                 sample_weight_filt_t = sample_weight_filt[w == 1]
                 self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
@@ -333,19 +325,17 @@ def estimate_ate(
         if pretrain:
             te = self.predict(X, p)
         else:
-            if (
-                treatment_np is None
-                or not len(treatment_np)
-                or y_np is None
-                or not len(y_np)
-            ):
-                if not len(treatment) or not len(y):
-                    raise ValueError(
-                        "treatment and y must be provided when pretrain=False"
-                    )
-                te = self.fit_predict(
-                    X, treatment, y, p, sample_weight, return_ci=False
-                )
+            if treatment is None or y is None:
+                raise ValueError("treatment and y must be provided when pretrain=False")
+
+            te = self.fit_predict(
+                X,
+                treatment,
+                y,
+                p,
+                sample_weight,
+                return_ci=False,
+            )
 
         ate = np.zeros(self.t_groups.shape[0])
         ate_lb = np.zeros(self.t_groups.shape[0])
@@ -577,9 +567,6 @@ def predict(self, X, p=None):
         """
         X = collect_if_lazy(X)
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
-        """Predict treatment effects."""
-        X = convert_pd_to_np(X)
-        te = np.zeros((X.shape[0], self.t_groups.shape[0]))
         for i, group in enumerate(self.t_groups):
             te[:, i] = self.models_tau[group].predict(X)
         return te

From 509a213514a58d956b1b99833f56a1f1ce1b844d Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Tue, 30 Jun 2026 23:48:27 +0530
Subject: [PATCH 14/15] fixing blockers

---
 causalml/inference/meta/drlearner.py |  63 ++++-----------
 causalml/inference/meta/rlearner.py  |  54 ++++---------
 causalml/inference/meta/slearner.py  |  16 +---
 causalml/inference/meta/tlearner.py  |  17 +---
 causalml/inference/meta/xlearner.py  | 111 +++++++--------------------
 5 files changed, 60 insertions(+), 201 deletions(-)

diff --git a/causalml/inference/meta/drlearner.py b/causalml/inference/meta/drlearner.py
index 942a56f3..5649f376 100644
--- a/causalml/inference/meta/drlearner.py
+++ b/causalml/inference/meta/drlearner.py
@@ -16,7 +16,6 @@
     filter_index,
     n_rows,
     to_numpy,
-    convert_pd_to_np,
 )
 from causalml.metrics import regression_metrics, classification_metrics
 from causalml.propensity import compute_propensity_score
@@ -51,33 +50,11 @@ def __init__(
             treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
             ate_alpha (float, optional): the confidence level alpha of the ATE estimate
             control_name (str or int, optional): name of control group
-        """
-        assert (learner is not None) or (
-            (control_outcome_learner is not None)
-            and (treatment_outcome_learner is not None)
-            and (treatment_effect_learner is not None)
-        )
 
-        self.model_mu_c = (
-            deepcopy(learner)
-            if control_outcome_learner is None
-            else control_outcome_learner
-        )
-        self.model_mu_t = (
-            deepcopy(learner)
-            if treatment_outcome_learner is None
-            else treatment_outcome_learner
-        )
-        self.model_tau = (
-            deepcopy(learner)
-            if treatment_effect_learner is None
-            else treatment_effect_learner
-        )
-        """
         Note: arguments are stored verbatim (scikit-learn convention) so that
-        ``get_params`` / ``clone`` work correctly. Model construction is deferred to ``fit()``.
-        Per the scikit-learn convention, ``__init__`` does not validate or raise —
-        validation happens in ``fit()``.
+        ``get_params`` / ``clone`` work correctly. Model construction is deferred
+        to ``fit()``. Per the scikit-learn convention, ``__init__`` does not
+        validate or raise — validation happens in ``fit()``.
         """
         # Store verbatim — no deepcopy, no logic (scikit-learn convention).
         self.learner = learner
@@ -104,14 +81,8 @@ def fit(self, X, treatment, y, p=None, seed=None):
                 single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
                 float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
             seed (int): random seed for cross-fitting
-
-        X = collect_if_lazy(X)
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
-            seed (int): random seed for cross-fitting
         """
+        X = collect_if_lazy(X)
         if (self.learner is None) and (
             (self.control_outcome_learner is None)
             or (self.treatment_outcome_learner is None)
@@ -122,7 +93,6 @@ def fit(self, X, treatment, y, p=None, seed=None):
                 "`treatment_outcome_learner`, and `treatment_effect_learner` "
                 "must be specified."
             )
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
@@ -305,8 +275,6 @@ def predict(
         X = collect_if_lazy(X)
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
-        yhat_cs = {}
-
         yhat_ts = {}
 
         yhat_c = np.r_[[model.predict(X) for model in self.models_mu_c]].mean(axis=0)
@@ -352,12 +320,15 @@ def fit_predict(
         verbose=True,
         seed=None,
     ):
-        """
+        """Fit the treatment effect and outcome models of the DR learner and predict treatment effects.
+
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
+                single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
             return_ci (bool): whether to return confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -451,10 +422,6 @@ def estimate_ate(
             p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
                 single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
                 float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
             bootstrap_ci (bool): whether run bootstrap for confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -613,7 +580,7 @@ def __init__(
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        """Predict treatment effects.
+        """Predict treatment effects (classifier variant — uses predict_proba for outcomes).
 
         Args:
             X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
@@ -637,10 +604,6 @@ def predict(
         X = collect_if_lazy(X)
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
-        yhat_cs = {}
-        """Predict treatment effects (classifier variant — uses predict_proba for outcomes)."""
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
-
         yhat_ts = {}
 
         yhat_c = np.r_[
diff --git a/causalml/inference/meta/rlearner.py b/causalml/inference/meta/rlearner.py
index f0aff405..3120038f 100644
--- a/causalml/inference/meta/rlearner.py
+++ b/causalml/inference/meta/rlearner.py
@@ -15,7 +15,6 @@
     to_numpy,
     get_xgboost_objective_metric,
     get_weighted_variance,
-    convert_pd_to_np,
 )
 from causalml.propensity import ElasticNetPropensityModel
 
@@ -59,9 +58,10 @@ def __init__(
                 processors
 
         Note: arguments are stored verbatim (scikit-learn convention) so that
-        ``get_params`` / ``clone`` work correctly. Model construction is deferred to ``fit()``.
-        Per the scikit-learn convention, ``__init__`` does not validate or raise —
-        validation of ``learner``/``outcome_learner``/``effect_learner`` happens in ``fit()``.
+        ``get_params`` / ``clone`` work correctly. Model construction is deferred
+        to ``fit()``. Per the scikit-learn convention, ``__init__`` does not
+        validate or raise — validation of ``learner``/``outcome_learner``/
+        ``effect_learner`` happens in ``fit()``.
         """
         # Store verbatim — no deepcopy, no logic (scikit-learn convention).
         self.learner = learner
@@ -91,13 +91,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
                 weight of each observation for `effect_learner`. If None, it assumes equal weight.
             verbose (bool, optional): whether to output progress logs
-
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
-            sample_weight (np.array or pd.Series, optional): sample weights for `effect_learner`.
-            verbose (bool, optional): whether to output progress logs
         """
         X = collect_if_lazy(X)
         if (self.learner is None) and (
@@ -109,7 +102,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             )
         if self.propensity_learner is None:
             raise ValueError("`propensity_learner` must be specified.")
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
@@ -172,11 +164,12 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             diff_t = y_filt[w == 1] - yhat_filt[w == 1]
             if sample_weight is not None:
                 sample_weight_filt = sample_weight[mask]
-
-                sample_weight_filt_c = sample_weight_filt[w == 0]
-                sample_weight_filt_t = sample_weight_filt[w == 1]
-                self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
-                self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
+                self.vars_c[group] = get_weighted_variance(
+                    diff_c, sample_weight_filt[w == 0]
+                )
+                self.vars_t[group] = get_weighted_variance(
+                    diff_t, sample_weight_filt[w == 1]
+                )
                 weight *= sample_weight_filt
             else:
                 self.vars_c[group] = diff_c.var()
@@ -232,11 +225,6 @@ def fit_predict(
                 float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
             sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
                 weight of each observation for `effect_learner`. If None, it assumes equal weight.
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
-            sample_weight (np.array or pd.Series, optional): sample weights
             return_ci (bool): whether to return confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -306,11 +294,6 @@ def estimate_ate(
                 float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
             sample_weight (np.array, pd.Series, or pl.Series, optional): an array of sample weights indicating the
                 weight of each observation for `effect_learner`. If None, it assumes equal weight.
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): treatment vector (needed when pretrain=False)
-            y (np.array or pd.Series): outcome vector (needed when pretrain=False)
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
-            sample_weight (np.array or pd.Series, optional): sample weights
             bootstrap_ci (bool): whether run bootstrap for confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -459,7 +442,7 @@ def __init__(
         )
 
     def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
-        """Fit the treatment effect and outcome models of the R learner.
+        """Fit the R-learner classifier (uses predict_proba for outcome estimates).
 
         Args:
             X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
@@ -474,8 +457,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
             verbose (bool, optional): whether to output progress logs
         """
         X = collect_if_lazy(X)
-        """Fit the R-learner classifier (uses predict_proba for outcome estimates)."""
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
@@ -535,10 +516,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
                 self.vars_t[group] = get_weighted_variance(
                     diff_t, sample_weight_filt[w == 1]
                 )
-                sample_weight_filt_c = sample_weight_filt[w == 0]
-                sample_weight_filt_t = sample_weight_filt[w == 1]
-                self.vars_c[group] = get_weighted_variance(diff_c, sample_weight_filt_c)
-                self.vars_t[group] = get_weighted_variance(diff_t, sample_weight_filt_t)
                 weight *= sample_weight_filt
             else:
                 self.vars_c[group] = diff_c.var()
@@ -621,9 +598,6 @@ def __init__(
         assert isinstance(random_state, int), "random_state should be int."
 
         # Store verbatim — no transformation, no XGBRegressor construction here.
-        # xgb_kwargs=None is stored as-is; BaseEstimator.get_params surfaces it
-        # correctly since it is a named parameter.  The or {} coalesce happens in
-        # fit() so that clone(XGBRRegressor()) still round-trips None → None.
         self.early_stopping = early_stopping
         self.test_size = test_size
         self.early_stopping_rounds = early_stopping_rounds
@@ -642,11 +616,12 @@ def __init__(
         )
 
     def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
-        """Fit the treatment effect and outcome models of the R learner.
+        """Fit using early-stopping XGBoost R-learner.
 
         Args:
             X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
                 A pl.LazyFrame is collected once at the start of this method.
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
             y (np.array, pd.Series, or pl.Series): an outcome vector
             p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in the
                 single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
@@ -661,9 +636,6 @@ def fit(self, X, treatment, y, p=None, sample_weight=None, verbose=True):
         y_np = to_numpy(y)
 
         # initialize equal sample weight if it's not provided, for simplicity purpose
-        """Fit using early-stopping XGBoost R-learner."""
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
-        check_treatment_vector(treatment, self.control_name)
         sample_weight = (
             to_numpy(sample_weight) if sample_weight is not None else np.ones(len(y_np))
         )
diff --git a/causalml/inference/meta/slearner.py b/causalml/inference/meta/slearner.py
index 7001a499..ce05ed9b 100644
--- a/causalml/inference/meta/slearner.py
+++ b/causalml/inference/meta/slearner.py
@@ -66,9 +66,8 @@ def __init__(self, learner=None, ate_alpha=0.05, control_name=0):
         """Initialize an S-learner.
 
         Args:
-            learner (optional): a model to estimate the treatment effect
             learner (optional): a model to estimate the treatment effect.
-                If None, a DummyRegressor is used.  The argument is stored
+                If None, a DummyRegressor is used. The argument is stored
                 verbatim so that ``get_params`` / ``clone`` work correctly
                 (scikit-learn convention).
             ate_alpha (float, optional): the confidence level alpha of the ATE estimate
@@ -130,10 +129,6 @@ def predict(
         yhat_cs = {}
         yhat_ts = {}
 
-        # Build the augmented arrays once; they are identical for every group.
-        # (Separate allocations avoid in-place mutation by learners like CatBoost
-        # that set the writeable flag to False on arrays passed to predict().)
-
         for group in self.t_groups:
             model = self.models[group]
 
@@ -141,7 +136,6 @@ def predict(
             # mutation, which fails when learners like CatBoost set the
             # writeable flag to False on arrays passed to predict().
             X_new_c = prepend_column(0.0, X)
-
             X_new_t = prepend_column(1.0, X)
             yhat_cs[group] = model.predict(X_new_c)
             yhat_ts[group] = model.predict(X_new_t)
@@ -382,15 +376,13 @@ def predict(
         yhat_cs = {}
         yhat_ts = {}
 
-        # Build the augmented arrays once; they are identical for every group.
-        X_new_c = np.hstack((np.zeros((X.shape[0], 1)), X))
-        X_new_t = np.hstack((np.ones((X.shape[0], 1)), X))
-
         for group in self.t_groups:
             model = self.models[group]
 
+            # Build separate frames for control and treatment to avoid in-place
+            # mutation, which fails when learners like CatBoost set the
+            # writeable flag to False on arrays passed to predict().
             X_new_c = prepend_column(0.0, X)
-
             X_new_t = prepend_column(1.0, X)
             yhat_cs[group] = model.predict_proba(X_new_c)[:, 1]
             yhat_ts[group] = model.predict_proba(X_new_t)[:, 1]
diff --git a/causalml/inference/meta/tlearner.py b/causalml/inference/meta/tlearner.py
index cd0487be..d9d555da 100644
--- a/causalml/inference/meta/tlearner.py
+++ b/causalml/inference/meta/tlearner.py
@@ -21,7 +21,6 @@
     filter_mask,
     n_rows,
     to_numpy,
-    convert_pd_to_np,
 )
 from causalml.metrics import regression_metrics, classification_metrics
 
@@ -104,7 +103,6 @@ def fit(
                 "Either `learner` or both `control_learner` and `treatment_learner` "
                 "must be specified."
             )
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         y_np = to_numpy(y)
@@ -113,10 +111,6 @@ def fit(
         self.t_groups.sort()
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
 
-        # model_c is trained on the control group, which is identical for every
-        # treatment group, so fit it once. Deepcopy from the unfitted template so
-        # re-calling fit() always starts from a clean state.
-
         # Resolve base models from stored constructor args (no templates needed).
         _control_learner = (
             self.control_learner
@@ -131,10 +125,11 @@ def fit(
 
         self.models_t = {group: deepcopy(_treatment_learner) for group in self.t_groups}
 
-        # model_c is trained on the control group, identical for every treatment group.
-        control_mask = treatment == self.control_name
+        # model_c is trained on the control group, which is identical for every
+        # treatment group, so fit it once.
+        control_mask = treatment_np == self.control_name
         self.model_c = deepcopy(_control_learner)
-        self.model_c.fit(X[control_mask], y[control_mask])
+        self.model_c.fit(filter_mask(X, control_mask), y_np[control_mask])
         # Expose as a shared-reference dict to preserve the public models_c API.
         self.models_c = {group: self.model_c for group in self.t_groups}
 
@@ -198,9 +193,6 @@ def predict(
             y (np.array, pd.Series, or pl.Series, optional): an outcome vector
             return_components (bool, optional): whether to return outcome for treatment and control seperately
             verbose (bool, optional): whether to output progress logs
-            return_ci (bool, optional): whether to return confidence intervals using
-                the stored bootstrap ensemble. Requires fit() to have been called
-                with store_bootstraps=True.
             return_ci (bool, optional): whether to return confidence intervals
                 using the stored bootstrap ensemble. Requires fit() to have been
                 called with store_bootstraps=True.
@@ -521,7 +513,6 @@ def predict(
                 yhat[w == 1] = yhat_ts[group][mask][w == 1]
 
                 logger.info("Error metrics for group {}".format(group))
-
                 classification_metrics(y_filt, yhat, w)
 
         te = np.zeros((n_rows(X), self.t_groups.shape[0]))
diff --git a/causalml/inference/meta/xlearner.py b/causalml/inference/meta/xlearner.py
index b93cb2fe..536b8425 100644
--- a/causalml/inference/meta/xlearner.py
+++ b/causalml/inference/meta/xlearner.py
@@ -11,7 +11,6 @@
     filter_mask,
     n_rows,
     to_numpy,
-    convert_pd_to_np,
 )
 from causalml.metrics import regression_metrics, classification_metrics
 
@@ -47,39 +46,11 @@ def __init__(
             treatment_effect_learner (optional): a model to estimate treatment effects in the treatment group
             ate_alpha (float, optional): the confidence level alpha of the ATE estimate
             control_name (str or int, optional): name of control group
-        """
-        assert (learner is not None) or (
-            (control_outcome_learner is not None)
-            and (treatment_outcome_learner is not None)
-            and (control_effect_learner is not None)
-            and (treatment_effect_learner is not None)
-        )
 
-        self.model_mu_c = (
-            deepcopy(learner)
-            if control_outcome_learner is None
-            else control_outcome_learner
-        )
-        self.model_mu_t = (
-            deepcopy(learner)
-            if treatment_outcome_learner is None
-            else treatment_outcome_learner
-        )
-        self.model_tau_c = (
-            deepcopy(learner)
-            if control_effect_learner is None
-            else control_effect_learner
-        )
-        self.model_tau_t = (
-            deepcopy(learner)
-            if treatment_effect_learner is None
-            else treatment_effect_learner
-        )
-        """
         Note: arguments are stored verbatim (scikit-learn convention) so that
-        ``get_params`` / ``clone`` work correctly. Model construction is deferred to ``fit()``.
-        Per the scikit-learn convention, ``__init__`` does not validate or raise —
-        validation happens in ``fit()``.
+        ``get_params`` / ``clone`` work correctly. Model construction is deferred
+        to ``fit()``. Per the scikit-learn convention, ``__init__`` does not
+        validate or raise — validation happens in ``fit()``.
         """
         # Store verbatim — no deepcopy, no logic (scikit-learn convention).
         self.learner = learner
@@ -90,8 +61,8 @@ def __init__(
         self.ate_alpha = ate_alpha
         self.control_name = control_name
         # Sentinel so estimate_ate(pretrain=True) raises a clean ValueError
-        # ("no propensity score, please call fit() first") instead of AttributeError
-        # when called before fit().
+        # ("no propensity score, please call fit() first") instead of
+        # AttributeError when called before fit().
         self.propensity = {}
 
     def fit(self, X, treatment, y, p=None):
@@ -119,7 +90,6 @@ def fit(self, X, treatment, y, p=None):
                 "`treatment_outcome_learner`, `control_effect_learner`, and "
                 "`treatment_effect_learner` must be specified."
             )
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
@@ -133,7 +103,7 @@ def fit(self, X, treatment, y, p=None):
 
         self._classes = {group: i for i, group in enumerate(self.t_groups)}
 
-        # Resolve base models from stored constructor args (no templates needed).
+        # Resolve base models from stored constructor args (scikit-learn convention).
         _control_outcome_learner = (
             self.control_outcome_learner
             if self.control_outcome_learner is not None
@@ -170,20 +140,11 @@ def fit(self, X, treatment, y, p=None):
         control_mask = treatment_np == self.control_name
         X_control = filter_mask(X, control_mask)
         y_control = to_numpy(filter_mask(y, control_mask))
-        self.model_mu_c = deepcopy(self.model_mu_c)
+        self.model_mu_c = deepcopy(_control_outcome_learner)
         self.model_mu_c.fit(X_control, y_control)
         self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
-        # var_c is a single scalar since control model is shared across groups
+        # var_c is a single scalar since the control model is shared across groups
         self.var_c = (y_control - self.model_mu_c.predict(X_control)).var()
-        # Keep vars_c dict for backward compat with estimate_ate
-        # model_mu_c is trained on control data, identical for every treatment group.
-        control_mask = treatment == self.control_name
-        self.model_mu_c = deepcopy(_control_outcome_learner)
-        self.model_mu_c.fit(X[control_mask], y[control_mask])
-        self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
-
-        y_control_pred = self.model_mu_c.predict(X[control_mask])
-        self.var_c = (y[control_mask] - y_control_pred).var()
         self.vars_c = {group: self.var_c for group in self.t_groups}
 
         for group in self.t_groups:
@@ -210,6 +171,7 @@ def fit(self, X, treatment, y, p=None):
             d_t = y_filt_np[w == 1] - self.model_mu_c.predict(X_filt_t)
             self.models_tau_c[group].fit(X_filt_c, d_c)
             self.models_tau_t[group].fit(X_filt_t, d_t)
+
         return self
 
     def predict(
@@ -225,10 +187,6 @@ def predict(
             p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
                 the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
                 float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series, optional): a treatment vector
-            y (np.array or pd.Series, optional): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
             return_components (bool, optional): whether to return outcome for treatment and control seperately
             verbose (bool, optional): whether to output progress logs
         Returns:
@@ -249,11 +207,6 @@ def predict(
         dhat_cs = {}
         dhat_ts = {}
 
-        yhat_c_verbose = None
-        if (y is not None) and (treatment is not None) and verbose:
-            control_mask = treatment == self.control_name
-            yhat_c_verbose = self.model_mu_c.predict(X[control_mask])
-
         for i, group in enumerate(self.t_groups):
             dhat_cs[group] = self.models_tau_c[group].predict(X)
             dhat_ts[group] = self.models_tau_t[group].predict(X)
@@ -299,14 +252,15 @@ def fit_predict(
         return_components=False,
         verbose=True,
     ):
-        """
-        Fit the X-learner and predict treatment effects.
+        """Fit the X-learner and predict treatment effects.
 
         Args:
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
+            X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix
+            treatment (np.array, pd.Series, or pl.Series): a treatment vector
+            y (np.array, pd.Series, or pl.Series): an outcome vector
+            p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
+                the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
+                float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
             return_ci (bool): whether to return confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -382,10 +336,6 @@ def estimate_ate(
             p (np.ndarray, pd.Series, pl.Series, or dict, optional): an array of propensity scores of float (0,1) in
                 the single-treatment case; or, a dictionary of treatment groups that map to propensity vectors of
                 float (0,1); if None will run ElasticNetPropensityModel() to generate the propensity scores.
-            X (np.matrix or np.array or pd.Dataframe): a feature matrix
-            treatment (np.array or pd.Series): a treatment vector
-            y (np.array or pd.Series): an outcome vector
-            p (np.ndarray or pd.Series or dict, optional): propensity scores
             bootstrap_ci (bool): whether run bootstrap for confidence intervals
             n_bootstraps (int): number of bootstrap iterations
             bootstrap_size (int): number of samples per bootstrap
@@ -568,7 +518,8 @@ def fit(self, X, treatment, y, p=None):
             or (self.treatment_effect_learner is None)
         ):
             raise ValueError(
-                "Either `outcome_learner` and `effect_learner`, or all four specialized learners must be specified."
+                "Either `outcome_learner` and `effect_learner`, or all four "
+                "specialized learners must be specified."
             )
 
         _control_outcome_learner = (
@@ -576,25 +527,22 @@ def fit(self, X, treatment, y, p=None):
             if self.control_outcome_learner is None
             else deepcopy(self.control_outcome_learner)
         )
-
         _treatment_outcome_learner = (
             deepcopy(self.outcome_learner)
             if self.treatment_outcome_learner is None
             else deepcopy(self.treatment_outcome_learner)
         )
-
         _control_effect_learner = (
             deepcopy(self.effect_learner)
             if self.control_effect_learner is None
             else deepcopy(self.control_effect_learner)
         )
-
         _treatment_effect_learner = (
             deepcopy(self.effect_learner)
             if self.treatment_effect_learner is None
             else deepcopy(self.treatment_effect_learner)
         )
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
+
         check_treatment_vector(treatment, self.control_name)
         treatment_np = to_numpy(treatment)
         self.t_groups = np.unique(treatment_np[treatment_np != self.control_name])
@@ -620,13 +568,13 @@ def fit(self, X, treatment, y, p=None):
         self.vars_t = {}
 
         # model_mu_c is trained on control only (identical across groups) — fit once.
-        control_mask = treatment == self.control_name
+        control_mask = treatment_np == self.control_name
+        X_control = filter_mask(X, control_mask)
+        y_control = to_numpy(filter_mask(y, control_mask))
         self.model_mu_c = deepcopy(_control_outcome_learner)
-        self.model_mu_c.fit(X[control_mask], y[control_mask])
+        self.model_mu_c.fit(X_control, y_control)
         self.models_mu_c = {group: self.model_mu_c for group in self.t_groups}
-
-        y_control_pred = self.model_mu_c.predict_proba(X[control_mask])[:, 1]
-        self.var_c = (y[control_mask] - y_control_pred).var()
+        self.var_c = (y_control - self.model_mu_c.predict_proba(X_control)[:, 1]).var()
         self.vars_c = {group: self.var_c for group in self.t_groups}
 
         for group in self.t_groups:
@@ -649,7 +597,7 @@ def fit(self, X, treatment, y, p=None):
             ).var()
             self.vars_t[group] = var_t
 
-            # Train treatment models
+            # Train treatment effect models
             d_c = (
                 self.models_mu_t[group].predict_proba(X_filt_c)[:, 1]
                 - y_filt_np[w == 0]
@@ -666,7 +614,7 @@ def fit(self, X, treatment, y, p=None):
     def predict(
         self, X, treatment=None, y=None, p=None, return_components=False, verbose=True
     ):
-        """Predict treatment effects.
+        """Predict treatment effects (classifier variant — uses predict_proba).
 
         Args:
             X (np.matrix, np.array, pd.DataFrame, pl.DataFrame, or pl.LazyFrame): a feature matrix.
@@ -682,8 +630,6 @@ def predict(
             (numpy.ndarray): Predictions of treatment effects.
         """
         X = collect_if_lazy(X)
-        """Predict treatment effects (classifier variant — uses predict_proba)."""
-        X, treatment, y = convert_pd_to_np(X, treatment, y)
 
         if p is None:
             logger.info("Generating propensity score")
@@ -698,11 +644,6 @@ def predict(
         dhat_cs = {}
         dhat_ts = {}
 
-        yhat_c_verbose = None
-        if (y is not None) and (treatment is not None) and verbose:
-            control_mask = treatment == self.control_name
-            yhat_c_verbose = self.model_mu_c.predict_proba(X[control_mask])[:, 1]
-
         for i, group in enumerate(self.t_groups):
             dhat_cs[group] = self.models_tau_c[group].predict(X)
             dhat_ts[group] = self.models_tau_t[group].predict(X)

From 561ea3e07fccbabfe581155beab7ca543149144a Mon Sep 17 00:00:00 2001
From: Aman Srivastava <amansri345@gmail.com>
Date: Wed, 1 Jul 2026 12:56:28 +0530
Subject: [PATCH 15/15] _fit_bootstrap_clone fix

---
 causalml/inference/meta/base.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/causalml/inference/meta/base.py b/causalml/inference/meta/base.py
index 8c191292..18d6f6f2 100644
--- a/causalml/inference/meta/base.py
+++ b/causalml/inference/meta/base.py
@@ -36,8 +36,9 @@ def _fit_bootstrap_clone(learner_template, X, treatment, y, p, seed, bootstrap_s
         A fitted clone of learner_template trained on a bootstrap sample.
     """
     rng = np.random.RandomState(seed)
-    idxs = rng.choice(np.arange(X.shape[0]), size=bootstrap_size)
-    X_b = X[idxs]
+    idxs = rng.choice(np.arange(n_rows(X)), size=bootstrap_size)
+
+    X_b = filter_index(X, idxs)
     treatment_b = treatment[idxs]
     y_b = y[idxs]
     p_b = {group: _p[idxs] for group, _p in p.items()} if p is not None else None