How to use the econml.cate_estimator.BaseCateEstimator function in econml

    @BaseCateEstimator._wrap_fit
    def fit(self, Y, T, X, W=None, inference=None):
        """Build an orthogonal random forest from a training set (Y, T, X, W).

        Parameters
        ----------
        Y : array-like, shape (n, )
            Outcome for the treatment policy.

        T : array-like, shape (n, d_t)
            Treatment policy.

        X : array-like, shape (n, d_x)
            Feature vector that captures heterogeneity.

        W : array-like, shape (n, d_w) or None (default=None)
            High-dimensional controls.

For binary treatments, it returns the same as `effect`.

        Parameters
        ----------
        X : matrix, shape (m × dₓ)
            Matrix of features for each sample.

        Returns
        -------
        τ_hat : array-like, shape (m, )
            Matrix of heterogeneous treatment effects for each sample.
        """
        return self.effect(X)


class SLearner(BaseCateEstimator):
    """Conditional mean regression estimator where the treatment assignment is taken as a feature in the ML model.

    Parameters
    ----------
    overall_model : outcome estimator for all units
        Model will be trained on X|T where '|' denotes concatenation.
        Must implement `fit` and `predict` methods.

    """

    def __init__(self, overall_model):
        self.overall_model = clone(overall_model, safe=False)
        super().__init__()

    @BaseCateEstimator._wrap_fit
    def fit(self, Y, T, X, inference=None):

    @BaseCateEstimator._wrap_fit
    def fit(self, Y, T, X, inference=None):
        """Build an instance of DomainAdaptationLearner.

        Parameters
        ----------
        Y : array-like, shape (n, ) or (n, d_y)
            Outcome(s) for the treatment policy.

        T : array-like, shape (n, ) or (n, 1)
            Treatment policy. Only binary treatments are accepted as input.
            T will be flattened if shape is (n, 1).

        X : array-like, shape (n, d_x)
            Feature vector that captures heterogeneity.

        inference: string, `Inference` instance, or None

    @BaseCateEstimator._wrap_fit
    def fit(self, Y, T, X, inference=None):
        """Build an instance of SLearner.

        Parameters
        ----------
        Y : array-like, shape (n, ) or (n, d_y)
            Outcome(s) for the treatment policy.

        T : array-like, shape (n, ) or (n, 1)
            Treatment policy. Only binary treatments are accepted as input.
            T will be flattened if shape is (n, 1).

        X : array-like, shape (n, d_x)
            Feature vector that captures heterogeneity.

        inference: string, `Inference` instance, or None

    @BaseCateEstimator._wrap_fit
    def fit(self, Y, T, X=None, W=None, Z=None, sample_weight=None, sample_var=None, *, inference=None):
        """
        Estimate the counterfactual model from data, i.e. estimates function :math:`\\theta(\\cdot)`.

        Parameters
        ----------
        Y: (n, d_y) matrix or vector of length n
            Outcomes for each sample
        T: (n, d_t) matrix or vector of length n
            Treatments for each sample
        X: optional (n, d_x) matrix or None (Default=None)
            Features for each sample
        W: optional (n, d_w) matrix or None (Default=None)
            Controls for each sample
        Z: optional (n, d_z) matrix or None (Default=None)
            Instruments for each sample

else:
                    columns.append(np.zeros((n, (self._degree + 1) * ncols)))
        return reshape(np.hstack(columns), (n,) + (ncols,) * self._shift + (-1,))


def _add_ones(arr):
    """Add a column of ones to the front of an array."""
    return np.hstack([np.ones((shape(arr)[0], 1)), arr])


def _add_zeros(arr):
    """Add a column of zeros to the front of an array."""
    return np.hstack([np.zeros((shape(arr)[0], 1)), arr])


class NonparametricTwoStageLeastSquares(BaseCateEstimator):
    """
    Non-parametric instrumental variables estimator.

    Supports the use of arbitrary featurizers for the features, treatments, and instruments.

    Parameters
    ----------
    t_featurizer: transformer
        Featurizer used to transform the treatments

    x_featurizer: transformer
        Featurizer used to transform the raw features

    z_featurizer: transformer
        Featurizer used to transform the instruments

Base treatments for each sample
        X: optional (m × dₓ) matrix
            Features for each sample

        Returns
        -------
        grad_tau: (m × d_y × dₜ) array
            Heterogeneous marginal effects on each outcome for each sample
            Note that when Y or T is a vector rather than a 2-dimensional array,
            the corresponding singleton dimensions in the output will be collapsed
            (e.g. if both are vectors, then the output of this method will also be a vector)
        """
        pass


class LinearCateEstimator(BaseCateEstimator):
    """Base class for all CATE estimators with linear treatment effects in this package."""

    @abc.abstractmethod
    def const_marginal_effect(self, X=None):
        """
        Calculate the constant marginal CATE θ(·).

        The marginal effect is conditional on a vector of
        features on a set of m test samples {Xᵢ}.

        Parameters
        ----------
        X: optional (m × dₓ) matrix
            Features for each sample

        Returns

For binary treatments, it returns the same as `effect`.

        Parameters
        ----------
        X : matrix, shape (m × dₓ)
            Matrix of features for each sample.

        Returns
        -------
        τ_hat : array-like, shape (m, )
            Matrix of heterogeneous treatment effects for each sample.
        """
        return self.effect(X)


class XLearner(BaseCateEstimator):
    """Meta-algorithm proposed by Kunzel et al. that performs best in settings
       where the number of units in one treatment arm is much larger than in the other.

    Parameters
    ----------
    controls_model : outcome estimator for control units
        Must implement `fit` and `predict` methods.

    treated_model : outcome estimator for treated units
        Must implement `fit` and `predict` methods.

    cate_controls_model : estimator for pseudo-treatment effects on the controls
        Must implement `fit` and `predict` methods.

    cate_treated_model : estimator for pseudo-treatment effects on the treated
        Must implement `fit` and `predict` methods.

For more details on these CATE methods, see 
(Künzel S., Sekhon J., Bickel P., Yu B.) on Arxiv.
"""

import numpy as np
import warnings
from .cate_estimator import BaseCateEstimator
from sklearn import clone
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.utils import check_array, check_X_y
from .utilities import check_inputs


class TLearner(BaseCateEstimator):
    """Conditional mean regression estimator.

    Parameters
    ----------
    controls_model : outcome estimator for control units
        Must implement `fit` and `predict` methods.

    treated_model : outcome estimator for treated units
        Must implement `fit` and `predict` methods.

    """

    def __init__(self, controls_model, treated_model):
        self.controls_model = clone(controls_model, safe=False)
        self.treated_model = clone(treated_model, safe=False)
        super().__init__()

For binary treatments, it returns the same as `effect`.

        Parameters
        ----------
        X : matrix, shape (m × dₓ)
            Matrix of features for each sample.

        Returns
        -------
        τ_hat : array-like, shape (m, )
            Matrix of heterogeneous treatment effects for each sample.
        """
        return self.effect(X)


class DomainAdaptationLearner(BaseCateEstimator):
    """Meta-algorithm that uses domain adaptation techniques to account for
       covariate shift (selection bias) between the treatment arms.

    Parameters
    ----------
    controls_model : outcome estimator for control units
        Must implement `fit` and `predict` methods.
        The `fit` method must accept the `sample_weight` parameter.

    treated_model : outcome estimator for treated units
        Must implement `fit` and `predict` methods.
        The `fit` method must accept the `sample_weight` parameter.

    overall_model : estimator for pseudo-treatment effects
        Must implement `fit` and `predict` methods.