ENH: Pass fit_params through train methods

Signed-off-by: Atharva Kelkar <[email protected]>

Author: atharva-novi

econml/dml/_rlearner.py

@@ -50,12 +50,12 @@ def __init__(self, model_y: ModelSelector, model_t: ModelSelector):
         self._model_y = model_y
         self._model_t = model_t
 
-    def train(self, is_selecting, folds, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None):
+    def train(self, is_selecting, folds, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None, **fit_params):
         assert Z is None, "Cannot accept instrument!"
         self._model_t.train(is_selecting, folds, X, W, T, **
-                            filter_none_kwargs(sample_weight=sample_weight, groups=groups))
+                            filter_none_kwargs(sample_weight=sample_weight, groups=groups, **fit_params))
         self._model_y.train(is_selecting, folds, X, W, Y, **
-                            filter_none_kwargs(sample_weight=sample_weight, groups=groups))
+                            filter_none_kwargs(sample_weight=sample_weight, groups=groups, **fit_params))
         return self
 
     def score(self, Y, T, X=None, W=None, Z=None, sample_weight=None, groups=None,

econml/dml/dml.py

@@ -97,7 +97,7 @@ def __init__(self, model: SingleModelSelector, discrete_target):
         self._model = clone(model, safe=False)
         self._discrete_target = discrete_target
 
-    def train(self, is_selecting, folds, X, W, Target, sample_weight=None, groups=None):
+    def train(self, is_selecting, folds, X, W, Target, sample_weight=None, groups=None, **fit_params):
         if self._discrete_target:
             # In this case, the Target is the one-hot-encoding of the treatment variable
             # We need to go back to the label representation of the one-hot so as to call
@@ -108,7 +108,7 @@ def train(self, is_selecting, folds, X, W, Target, sample_weight=None, groups=No
             Target = inverse_onehot(Target)
 
         self._model.train(is_selecting, folds, _combine(X, W, Target.shape[0]), Target,
-                          **filter_none_kwargs(groups=groups, sample_weight=sample_weight))
+                          **filter_none_kwargs(groups=groups, sample_weight=sample_weight, **fit_params))
         return self
 
     @property