Fix docstrings of fit() methods

Author: okiner-3

dte_adj/local.py

@@ -40,11 +40,11 @@ def fit(
         Train the SimpleLocalDistributionEstimator.
 
         Args:
-            covariates (ArrayLike): Pre-treatment covariates.
-            treatment_arms (ArrayLike): Treatment assignment variable (Z).
-            treatment_indicator (ArrayLike): Treatment indicator variable (D).
-            outcomes (ArrayLike): Scalar-valued observed outcome.
-            strata (ArrayLike): Stratum indicators.
+            covariates: Pre-treatment covariates.
+            treatment_arms: Treatment assignment variable (Z).
+            treatment_indicator: Treatment indicator variable (D).
+            outcomes: Scalar-valued observed outcome.
+            strata: Stratum indicators.
 
         Returns:
             SimpleLocalDistributionEstimator: The fitted estimator.
@@ -209,11 +209,11 @@ def fit(
         Train the AdjustedLocalDistributionEstimator.
 
         Args:
-            covariates (ArrayLike): Pre-treatment covariates.
-            treatment_arms (ArrayLike): Treatment assignment variable (Z).
-            treatment_indicator (ArrayLike): Treatment indicator variable (D).
-            outcomes (ArrayLike): Scalar-valued observed outcome.
-            strata (ArrayLike): Stratum indicators.
+            covariates: Pre-treatment covariates.
+            treatment_arms: Treatment assignment variable (Z).
+            treatment_indicator: Treatment indicator variable (D).
+            outcomes: Scalar-valued observed outcome.
+            strata: Stratum indicators.
 
         Returns:
             AdjustedLocalDistributionEstimator: The fitted estimator.

dte_adj/simple.py

@@ -54,9 +54,9 @@ def fit(
         Set parameters.
 
         Args:
-            covariates (ArrayLike): Pre-treatment covariates.
-            treatment_arms (ArrayLike): The index of the treatment arm.
-            outcomes (ArrayLike): Scalar-valued observed outcome.
+            covariates: Pre-treatment covariates.
+            treatment_arms: The index of the treatment arm.
+            outcomes: Scalar-valued observed outcome.
 
         Returns:
             SimpleDistributionEstimator: The fitted estimator.
@@ -118,9 +118,9 @@ def fit(
         Set parameters.
 
         Args:
-            covariates (ArrayLike): Pre-treatment covariates.
-            treatment_arms (ArrayLike): The index of the treatment arm.
-            outcomes (ArrayLike): Scalar-valued observed outcome.
+            covariates: Pre-treatment covariates.
+            treatment_arms: The index of the treatment arm.
+            outcomes: Scalar-valued observed outcome.
 
         Returns:
             AdjustedDistributionEstimator: The fitted estimator.

dte_adj/stratified.py

@@ -21,10 +21,10 @@ def fit(
         Train the DistributionEstimatorBase.
 
         Args:
-            covariates (ArrayLike): Pre-treatment covariates.
-            treatment_arms (ArrayLike): The index of the treatment arm.
-            outcomes (ArrayLike): Scalar-valued observed outcome.
-            strata (ArrayLike): Stratum indicators.
+            covariates: Pre-treatment covariates.
+            treatment_arms: The index of the treatment arm.
+            outcomes: Scalar-valued observed outcome.
+            strata: Stratum indicators.
 
         Returns:
             DistributionEstimatorBase: The fitted estimator.
@@ -186,10 +186,10 @@ def fit(
         Train the DistributionEstimatorBase.
 
         Args:
-            covariates (ArrayLike): Pre-treatment covariates.
-            treatment_arms (ArrayLike): The index of the treatment arm.
-            outcomes (ArrayLike): Scalar-valued observed outcome.
-            strata (ArrayLike): Stratum indicators.
+            covariates: Pre-treatment covariates.
+            treatment_arms: The index of the treatment arm.
+            outcomes: Scalar-valued observed outcome.
+            strata: Stratum indicators.
 
         Returns:
             DistributionEstimatorBase: The fitted estimator.