Conditionally .squeeze() probability distribution, handles cases of…

… 1-D probability distributions
infer-actively · Jul 27, 2022 · 584caf5 · 584caf5
1 parent 21b0a52
commit 584caf5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/pymdp/utils.py b/pymdp/utils.py
@@ -15,7 +15,8 @@
 EPS_VAL = 1e-16 # global constant for use in norm_dist()
 
 def sample(probabilities):
-    sample_onehot = np.random.multinomial(1, probabilities.squeeze())
+    probabilities = probabilities.squeeze() if len(probabilities) > 1 else probabilities
+    sample_onehot = np.random.multinomial(1, probabilities)
     return np.where(sample_onehot == 1)[0][0]
 
 def sample_obj_array(arr):