Implement calc_expectation
@@ -1337,6 +1337,16 @@ def calc_expectation(dstn, func=lambda x: x, *args):
"""
N = dstn.dim()
+ if isinstance(N, tuple):
+
+ f_query = list(
+ map(func, [dstn.X[...,i] for i in range(len(dstn.pmf))])
+ )
+ f_query = np.stack(f_query, axis = -1)
+
+ f_exp = np.dot(f_query, np.vstack(dstn.pmf))
+
+ else:
dstn_array = np.column_stack(dstn.X)
if N > 1:
