changed stable computation of 1/x

infer-actively · Jul 12, 2024 · 6d5b7f9 · 6d5b7f9
1 parent 799581a
commit 6d5b7f9
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/pymdp/jax/maths.py b/pymdp/jax/maths.py
@@ -134,9 +134,8 @@ def spm_wnorm(A):
     Returns Expectation of logarithm of Dirichlet parameters over a set of 
     Categorical distributions, stored in the columns of A.
     """
-    A = jnp.clip(A, min=MINVAL)
     norm = 1. / A.sum(axis=0)
-    avg = 1. / A
+    avg = 1. / (A + MINVAL)
     wA = norm - avg
     return wA