Fix wrong ZeroSumNormal logp expression

Author: lucianopaz

pymc/distributions/multivariate.py

@@ -2781,6 +2781,7 @@ def zerosumnormal_logp(op, values, normal_dist, sigma, support_shape, **kwargs):
     (value,) = values
     shape = value.shape
     n_zerosum_axes = op.ndim_supp
+    *_, sigma = normal_dist.owner.inputs
 
     _deg_free_support_shape = pt.inc_subtensor(shape[-n_zerosum_axes:], -1)
     _full_size = pm.floatX(pt.prod(shape))
@@ -2792,7 +2793,8 @@ def zerosumnormal_logp(op, values, normal_dist, sigma, support_shape, **kwargs):
     ]
 
     out = pt.sum(
-        pm.logp(normal_dist, value) * _degrees_of_freedom / _full_size,
+        -0.5 * pt.pow(value / sigma, 2)
+        - (pt.log(pt.sqrt(2.0 * np.pi)) + pt.log(sigma)) * _degrees_of_freedom / _full_size,
         axis=tuple(np.arange(-n_zerosum_axes, 0)),
     )
 

tests/distributions/test_multivariate.py

@@ -1596,7 +1596,7 @@ def test_zsn_variance(self, sigma, n):
             (5, 3, None, [-1]),
             (2, 6, None, [-1]),
             (5, (7, 3), None, [-1]),
-            (5, (2, 7, 3), 2, [1, 2]),
+            (5, (2, 7, 3), 2, [-2, -1]),
         ],
     )
     def test_zsn_logp(self, sigma, shape, n_zerosum_axes, mvn_axes):
@@ -1629,8 +1629,9 @@ def logp_norm(value, sigma, axes):
             return np.where(inds, np.sum(-psdet - exp, axis=-1), -np.inf)
 
         zsn_dist = pm.ZeroSumNormal.dist(sigma=sigma, shape=shape, n_zerosum_axes=n_zerosum_axes)
-        zsn_logp = pm.logp(zsn_dist, value=np.zeros(shape)).eval()
-        mvn_logp = logp_norm(value=np.zeros(shape), sigma=sigma, axes=mvn_axes)
+        zsn_draws = pm.draw(zsn_dist, 100)
+        zsn_logp = pm.logp(zsn_dist, value=zsn_draws).eval()
+        mvn_logp = logp_norm(value=zsn_draws, sigma=sigma, axes=mvn_axes)
 
         np.testing.assert_allclose(zsn_logp, mvn_logp)