diff --git a/HARK/ConsumptionSaving/ConsIndShockModel.py b/HARK/ConsumptionSaving/ConsIndShockModel.py
index 442809db9..b3f5dcdc8 100644
--- a/HARK/ConsumptionSaving/ConsIndShockModel.py
+++ b/HARK/ConsumptionSaving/ConsIndShockModel.py
@@ -793,9 +793,6 @@ def vp_next(shocks, a_nrm):
             )
         )
 
-        if EndOfPrdvP.shape[0] == EndOfPrdvP.size:
-            EndOfPrdvP = EndOfPrdvP.flatten()
-
         return EndOfPrdvP 
 
     def getPointsForInterpolation(self, EndOfPrdvP, aNrmNow):
diff --git a/HARK/distribution.py b/HARK/distribution.py
index 6a4d2bb1b..fb29e7236 100644
--- a/HARK/distribution.py
+++ b/HARK/distribution.py
@@ -1011,7 +1011,7 @@ def calcExpectation(dstn,func=lambda x : x,*args):
         The N-valued distribution over which the function is to be evaluated.
     func : function
         The function to be evaluated.
-        This function should take a 1D array of size N.
+        This function should take an array of size N x M.
         It may also take other arguments *args
         Please see numpy.apply_along_axis() for guidance on
         design of func.
@@ -1049,6 +1049,8 @@ def calcExpectation(dstn,func=lambda x : x,*args):
     # a hack.
     if f_exp.size == 1:
         f_exp = f_exp.flat[0]
+    elif f_exp.shape[0] == f_exp.size:
+        f_exp = f_exp.flatten()
 
     return f_exp
 
diff --git a/HARK/tests/test_distribution.py b/HARK/tests/test_distribution.py
index cf0effe50..3c32de9d6 100644
--- a/HARK/tests/test_distribution.py
+++ b/HARK/tests/test_distribution.py
@@ -94,7 +94,7 @@ def test_calcExpectation(self):
         )
 
         self.assertAlmostEqual(
-            ce9[3][0],
+            ce9[3],
             9.518015322143837
         )