Add separate function for check_y_params

Author: alexander-pv

causalml/inference/tree/causal/_tree.py

@@ -31,6 +31,19 @@
 }
 
 
+def get_check_y_params() -> dict:
+    """
+    Prepares flags depending on the scikit-learn version.
+
+    Returns: check_y_params
+    """
+    if Version(sklearn_version) >= Version("1.6"):
+        check_y_params = dict(ensure_2d=False, dtype=None, ensure_all_finite=False)
+    else:
+        check_y_params = dict(ensure_2d=False, dtype=None, force_all_finite=False)
+    return check_y_params
+
+
 class BaseCausalDecisionTree(BaseDecisionTree):
     """
     Modified base class BaseDecisionTree for causal trees
@@ -66,14 +79,7 @@ def fit(
             # Need to validate separately here.
             # We can't pass multi_ouput=True because that would allow y to be csr.
             check_X_params = dict(dtype=DTYPE, accept_sparse="csc")
-            if Version(sklearn_version) >= Version("1.6"):
-                check_y_params = dict(
-                    ensure_2d=False, dtype=None, ensure_all_finite=False
-                )
-            else:
-                check_y_params = dict(
-                    ensure_2d=False, dtype=None, force_all_finite=False
-                )
+            check_y_params = get_check_y_params()
             X, y = validate_data(
                 self, X, y, validate_separately=(check_X_params, check_y_params)
             )

causalml/inference/tree/causal/causalforest.py

@@ -19,6 +19,7 @@
 from sklearn.ensemble._forest import _generate_sample_indices, _get_n_samples_bootstrap
 
 from .causaltree import CausalTreeRegressor
+from ._tree import get_check_y_params
 
 try:
     from packaging.version import parse as Version
@@ -274,7 +275,7 @@ def _fit(
         if issparse(y):
             raise ValueError("sparse multilabel-indicator for y is not supported.")
         check_X_params = dict(dtype=DTYPE, accept_sparse="csc")
-        check_y_params = dict(ensure_2d=False, dtype=None, ensure_all_finite=False)
+        check_y_params = get_check_y_params()
         X, y = validate_data(
             self,
             X,