Handle single output and fix transform

Author: AlexAndorra

pymc/distributions/distribution.py

@@ -415,6 +415,14 @@ def dist(
 @_get_measurable_outputs.register(SymbolicRandomVariable)
 def _get_measurable_outputs_symbolic_random_variable(op, node):
     # This tells Aeppl that any non RandomType outputs are measurable
+
+    # Assume that if there is one default_output, that's the only one that is measurable
+    # In the rare case this is not what one wants, a specialized _get_measuarable_outputs
+    # can dispatch for a subclassed Op
+    if op.default_output is not None:
+        return [node.default_output()]
+
+    # Otherwise assume that any outputs that are not of RandomType are measurable
     return [out for out in node.outputs if not isinstance(out.type, RandomType)]
 
 

pymc/distributions/multivariate.py

@@ -2482,8 +2482,6 @@ def dist(cls, sigma=1, zerosum_axes=None, support_shape=None, **kwargs):
             ndim_supp=zerosum_axes,
         )
 
-        # print(f"{support_shape.eval() = }")
-
         if support_shape is None:
             if zerosum_axes > 0:
                 raise ValueError("You must specify shape or support_shape parameter")

pymc/distributions/transforms.py

@@ -291,7 +291,7 @@ def __init__(self, zerosum_axes):
             By default (``zerosum_axes=[-1]``), the sum-to-zero constraint is imposed
             on the last axis.
         """
-        self.zerosum_axes = zerosum_axes
+        self.zerosum_axes = tuple(int(axis) for axis in zerosum_axes)
 
     def forward(self, value, *rv_inputs):
         for axis in self.zerosum_axes: