Fix bug in supp_shape_from_ref_param_shape

Author: ricardoV94

pytensor/tensor/random/utils.py

@@ -323,7 +323,7 @@ def supp_shape_from_ref_param_shape(
         raise ValueError("ndim_supp must be greater than 0")
     if param_shapes is not None:
         ref_param = param_shapes[ref_param_idx]
-        return (ref_param[-ndim_supp],)
+        return tuple(ref_param[i] for i in range(-ndim_supp, 0))
     else:
         ref_param = dist_params[ref_param_idx]
         if ref_param.ndim < ndim_supp:

tests/tensor/random/test_utils.py

@@ -313,3 +313,11 @@ def test_supp_shape_from_ref_param_shape():
         ref_param_idx=1,
     )
     assert res == (3, 4)
+
+    res = supp_shape_from_ref_param_shape(
+        ndim_supp=2,
+        dist_params=(np.array([1, 2]), np.ones((2, 3, 4))),
+        param_shapes=((2,), (2, 3, 4)),
+        ref_param_idx=1,
+    )
+    assert res == (3, 4)