Allow Dim version of simple SymbolicRandomVariables

Author: ricardoV94

pymc/dims/distributions/scalar.py

@@ -14,15 +14,17 @@
 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 HalfStudentTRV, flat, halfflat
 from pymc.distributions.continuous import InverseGamma as RegularInverseGamma
-from pymc.distributions.continuous import flat, halfflat
 
 
 def _get_sigma_from_either_sigma_or_tau(*, sigma, tau):
@@ -90,6 +92,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

@@ -32,7 +32,7 @@
 from pytensor.graph.rewriting.basic import in2out
 from pytensor.graph.utils import MetaType
 from pytensor.tensor.basic import as_tensor_variable
-from pytensor.tensor.random.op import RandomVariable
+from pytensor.tensor.random.op import RandomVariable, RNGConsumerOp
 from pytensor.tensor.random.rewriting import local_subtensor_rv_lift
 from pytensor.tensor.random.utils import normalize_size_param
 from pytensor.tensor.rewriting.shape import ShapeFeature
@@ -207,7 +207,7 @@ def __get__(self, owner_self, owner_cls):
         return self.fget(owner_self if owner_self is not None else owner_cls)
 
 
-class SymbolicRandomVariable(MeasurableOp, OpFromGraph):
+class SymbolicRandomVariable(MeasurableOp, RNGConsumerOp, OpFromGraph):
     """Symbolic Random Variable.
 
     This is a subclasse of `OpFromGraph` which is used to encapsulate the symbolic
@@ -294,7 +294,10 @@ def default_output(cls_or_self) -> int | None:
     @staticmethod
     def get_input_output_type_idxs(
         extended_signature: str | None,
-    ) -> tuple[tuple[tuple[int], int | None, tuple[int]], tuple[tuple[int], tuple[int]]]:
+    ) -> tuple[
+        tuple[tuple[int, ...], int | None, tuple[int, ...]],
+        tuple[tuple[int, ...], tuple[int, ...]],
+    ]:
         """Parse extended_signature and return indexes for *[rng], [size] and parameters as well as outputs."""
         if extended_signature is None:
             raise ValueError("extended_signature must be provided")
@@ -367,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: