"Fixes linear DRIV shape inconsistency"

Signed-off-by: ssemov <ssemov@gmail.com>

econml/inference/_inference.py

@@ -13,7 +13,8 @@
 from ._bootstrap import BootstrapEstimator
 from ..utilities import (Summary, _safe_norm_ppf, broadcast_unit_treatments,
                          cross_product, inverse_onehot, ndim,
-                         parse_final_model_params, reshape_treatmentwise_effects, shape, filter_none_kwargs)
+                         parse_final_model_params, reshape_treatmentwise_effects, shape, filter_none_kwargs,
+                         reshape_outcomewise_effects)
 
 """Options for performing inference in estimators."""
 
@@ -281,8 +282,8 @@ def effect_inference(self, X, *, T0, T1):
         elif self.featurizer is not None:
             X = self.featurizer.transform(X)
         XT = cross_product(X, T1 - T0)
-        e_pred = self._predict(XT)
-        e_stderr = self._prediction_stderr(XT)
+        e_pred = reshape_outcomewise_effects(self._predict(XT), self._d_y)
+        e_stderr = reshape_outcomewise_effects(self._prediction_stderr(XT), self._d_y)
         d_y = self._d_y[0] if self._d_y else 1
 
         mean_XT = XT.mean(axis=0, keepdims=True)

econml/utilities.py

@@ -770,6 +770,27 @@ def reshape_treatmentwise_effects(A, d_t, d_y):
     else:
         return A
 
+def reshape_outcomewise_effects(A, d_y):
+    """
+    Given an effects matrix, reshape second dimension to be consistent with d_y[0].
+
+    Parameters
+    ----------
+    A : array
+        The effects array to be reshaped. It should have shape (m,) or (m, d_y).
+    d_y : tuple of int
+        Either () if Y was a vector, or a 1-tuple of the number of columns of Y if it was an array.
+
+    Returns
+    -------
+    A : array
+        The reshaped effects array with shape:
+         - (m, ) if d_y is () and Y is a vector,
+         - (m, d_y) if d_y is a 1-tuple and Y is an array.
+    """
+    if np.shape(A)[1:] == d_y:
+        return A
+    return A.reshape(-1, d_y[0])
 
 def einsum_sparse(subscripts, *arrs):
     """