Fix individual policy with no heterogeneity

Author: kbattocchi

econml/solutions/causal_analysis/_causal_analysis.py

@@ -1600,6 +1600,7 @@ def individualized_policy(self, Xtest, feature_index, *, n_rows=None, treatment_
 
         Xtest = result.X_transformer.transform(Xtest)
         if Xtest.shape[1] == 0:
+            x_rows = Xtest.shape[0]
             Xtest = None
 
         if result.feature_baseline is None:
@@ -1608,6 +1609,9 @@ def individualized_policy(self, Xtest, feature_index, *, n_rows=None, treatment_
         else:
             effect = result.estimator.const_marginal_effect_inference(Xtest)
 
+        if Xtest is None:  # we got a scalar effect although our original X may have had more rows
+            effect = effect._expand_outputs(x_rows)
+
         multi_y = (not self._vec_y) or self.classification
 
         if multi_y and result.feature_baseline is not None and np.ndim(treatment_costs) == 2:

econml/tests/test_causal_analysis.py

@@ -571,6 +571,7 @@ def test_empty_hinds(self):
                 eff = ca.local_causal_effect(X_df, alpha=0.05)
                 for ind in feat_inds:
                     pto = ca._policy_tree_output(X_df, ind)
+                    ca._individualized_policy_dict(X_df, ind)
 
     def test_can_serialize(self):
         import pickle