Allow Dim version of simple SymbolicRandomVariables

Author: ricardoV94

pymc/dims/distributions/scalar.py

@@ -14,14 +14,16 @@
 import pytensor.xtensor as ptx
 import pytensor.xtensor.random as pxr
 
+from pytensor.xtensor import as_xtensor
+
 from pymc.dims.distributions.core import (
     DimDistribution,
     PositiveDimDistribution,
     UnitDimDistribution,
 )
 from pymc.distributions.continuous import Beta as RegularBeta
 from pymc.distributions.continuous import Gamma as RegularGamma
-from pymc.distributions.continuous import flat, halfflat
+from pymc.distributions.continuous import HalfStudentTRV, flat, halfflat
 
 
 def _get_sigma_from_either_sigma_or_tau(*, sigma, tau):
@@ -89,6 +91,21 @@ def dist(cls, nu, mu=0, sigma=None, *, lam=None, **kwargs):
         return super().dist([nu, mu, sigma], **kwargs)
 
 
+class HalfStudentT(PositiveDimDistribution):
+    @classmethod
+    def dist(cls, nu, sigma=None, *, lam=None, **kwargs):
+        sigma = _get_sigma_from_either_sigma_or_tau(sigma=sigma, tau=lam)
+        return super().dist([nu, sigma], **kwargs)
+
+    @classmethod
+    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)
+        return xop(nu, sigma, core_dims=core_dims, extra_dims=extra_dims, rng=rng)
+
+
 class Cauchy(DimDistribution):
     xrv_op = pxr.cauchy
 

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: