ZeroSumNormal: Fix dims variant with batch sigma

Author: ricardoV94

pymc/dims/distributions/transforms.py

@@ -203,6 +203,4 @@ def backward(self, value, *rv_inputs):
         return value
 
     def log_jac_det(self, value, *rv_inputs):
-        # Use following once broadcast_like is implemented
-        # as_xtensor(0).broadcast_like(value, exclude=self.dims)`
-        return value.sum(self.dims) * 0
+        return as_xtensor(0.0).broadcast_like(value, exclude=self.dims)

pymc/dims/distributions/vector.py

@@ -229,7 +229,8 @@ def dist(cls, sigma=1.0, *, core_dims=None, dim_lengths, **kwargs):
             raise ValueError("ZeroSumNormal requires atleast 1 core_dims")
 
         support_dims = as_xtensor(
-            as_tensor([dim_lengths[core_dim] for core_dim in core_dims]), dims=("_",)
+            as_tensor([dim_lengths[core_dim] for core_dim in core_dims]),
+            dims=("__support_shape__",),
         )
         sigma = cls._as_xtensor(sigma)
 
@@ -238,16 +239,25 @@ def dist(cls, sigma=1.0, *, core_dims=None, dim_lengths, **kwargs):
         )
 
     @classmethod
-    def xrv_op(cls, sigma, support_dims, core_dims, extra_dims=None, rng=None):
-        sigma = cls._as_xtensor(sigma)
-        support_dims = as_xtensor(support_dims, dims=("_",))
-        support_shape = support_dims.values
-        core_rv = ZeroSumNormalRV.rv_op(sigma=sigma.values, support_shape=support_shape).owner.op
+    def xrv_op(cls, sigma, support_shape, core_dims, extra_dims=None, rng=None):
+        # ZeroSumNormal expects dummy dimensions on sigma for the support_shape
+        sigma = cls._as_xtensor(sigma).expand_dims(core_dims)
+        support_shape = as_xtensor(support_shape, dims=("__support_shape__",))
+        core_rv = ZeroSumNormalRV.rv_op(
+            sigma=sigma.values, support_shape=support_shape.values
+        ).owner.op
+        core_dims_map = tuple(range(1, len(core_dims) + 1))
         xop = pxr.as_xrv(
             core_rv,
-            core_inps_dims_map=[(), (0,)],
-            core_out_dims_map=tuple(range(1, len(core_dims) + 1)),
+            core_inps_dims_map=[core_dims_map, (0,)],
+            core_out_dims_map=core_dims_map,
         )
         # Dummy "_" core dim to absorb the support_shape vector
         # If ZeroSumNormal expected a scalar per support dim, this wouldn't be needed
-        return xop(sigma, support_dims, core_dims=("_", *core_dims), extra_dims=extra_dims, rng=rng)
+        return xop(
+            sigma,
+            support_shape,
+            core_dims=("__support_shape__", *core_dims),
+            extra_dims=extra_dims,
+            rng=rng,
+        )

pymc/distributions/distribution.py

@@ -392,7 +392,7 @@ def make_node(self, *inputs):
             )
             if size_arg_idx is not None and len(rng_arg_idxs) == 1:
                 new_size_type = normalize_size_param(inputs[size_arg_idx]).type
-                if not self.input_types[size_arg_idx].in_same_class(new_size_type):
+                if not self.input_types[size_arg_idx].is_super(new_size_type):
                     params = [inputs[idx] for idx in param_idxs]
                     size = inputs[size_arg_idx]
                     rng = inputs[rng_arg_idxs[0]]

tests/dims/distributions/test_vector.py

@@ -99,3 +99,23 @@ def test_zerosumnormal():
     # Logp is correct, but we have join(..., -1) and join(..., 1), that don't get canonicalized to the same
     # Should work once https://github.com/pymc-devs/pytensor/issues/1505 is fixed
     # assert_equivalent_logp_graph(model, reference_model)
+
+
+def test_zerosumnormal_batch_sigma():
+    coords = {"a": range(3), "b": range(5)}
+    sigma = np.array([1, 2, 3.0])
+    with Model(coords=coords) as model:
+        ZeroSumNormal(
+            "x",
+            sigma=as_xtensor(sigma, dims=("a",)),
+            core_dims=("b",),
+        )
+
+    with Model(coords=coords) as ref_model:
+        regular_distributions.ZeroSumNormal("x", sigma=sigma[:, None], dims=("a", "b"))
+
+    ip = model.initial_point()
+    np.testing.assert_allclose(
+        model.compile_logp(sum=False)(ip),
+        ref_model.compile_logp(sum=False)(ip),
+    )