Allow Dim version of simple SymbolicRandomVariables

Author: ricardoV94

pymc/dims/distributions/scalar.py

@@ -40,15 +40,15 @@ def _get_sigma_from_either_sigma_or_tau(*, sigma, tau):
 
 
 class Flat(DimDistribution):
-    xrv_op = pxr._as_xrv(flat)
+    xrv_op = pxr.as_xrv(flat)
 
     @classmethod
     def dist(cls, **kwargs):
         return super().dist([], **kwargs)
 
 
 class HalfFlat(PositiveDimDistribution):
-    xrv_op = pxr._as_xrv(halfflat, [], ())
+    xrv_op = pxr.as_xrv(halfflat, [], ())
 
     @classmethod
     def dist(cls, **kwargs):
@@ -102,7 +102,7 @@ def xrv_op(self, nu, sigma, core_dims=None, extra_dims=None, rng=None):
         nu = as_xtensor(nu)
         sigma = as_xtensor(sigma)
         core_rv = HalfStudentTRV.rv_op(nu=nu.values, sigma=sigma.values).owner.op
-        xop = pxr._as_xrv(core_rv)
+        xop = pxr.as_xrv(core_rv)
         return xop(nu, sigma, core_dims=core_dims, extra_dims=extra_dims, rng=rng)
 
 

pymc/distributions/distribution.py

@@ -370,8 +370,27 @@ def __init__(
 
         kwargs.setdefault("inline", True)
         kwargs.setdefault("strict", True)
+        # Many RVS have a size argument, even when this is `None` and is therefore unused
+        kwargs.setdefault("on_unused_input", "ignore")
         super().__init__(*args, **kwargs)
 
+    def make_node(self, *inputs):
+        # If we try to build the RV with a different size type (vector -> None or None -> vector)
+        # We need to rebuild the Op with new size type in the inner graph
+        if self.extended_signature is not None:
+            (rng_arg_idxs, size_arg_idx, param_idxs), _ = self.get_input_output_type_idxs(
+                self.extended_signature
+            )
+            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):
+                    params = [inputs[idx] for idx in param_idxs]
+                    size = inputs[size_arg_idx]
+                    rng = inputs[rng_arg_idxs[0]]
+                    return self.rebuild_rv(*params, size=size, rng=rng).owner
+
+        return super().make_node(*inputs)
+
     def update(self, node: Apply) -> dict[Variable, Variable]:
         """Symbolic update expression for input random state variables.
 

tests/dims/distributions/test_scalar.py

@@ -22,6 +22,7 @@
     HalfCauchy,
     HalfFlat,
     HalfNormal,
+    HalfStudentT,
     InverseGamma,
     Laplace,
     LogNormal,
@@ -119,6 +120,22 @@ def test_studentt():
     assert_equivalent_logp_graph(model, reference_model)
 
 
+def test_halfstudentt():
+    coords = {"a": range(3)}
+    with Model(coords=coords) as model:
+        HalfStudentT("x", nu=1, dims="a")
+        HalfStudentT("y", nu=1, sigma=3, dims="a")
+        HalfStudentT("z", nu=1, lam=3, dims="a")
+
+    with Model(coords=coords) as reference_model:
+        regular_distributions.HalfStudentT("x", nu=1, dims="a")
+        regular_distributions.HalfStudentT("y", nu=1, sigma=3, dims="a")
+        regular_distributions.HalfStudentT("z", nu=1, lam=3, dims="a")
+
+    assert_equivalent_random_graph(model, reference_model)
+    assert_equivalent_logp_graph(model, reference_model)
+
+
 def test_cauchy():
     coords = {"a": range(3)}
     with Model(coords=coords) as model: