Fail early in RandomVariable.make_node when size is incompatible with parameters dimensionality

Author: ricardoV94

pytensor/tensor/random/op.py

@@ -319,8 +319,23 @@ def make_node(self, rng, size, dtype, *dist_params):
                 "The type of rng should be an instance of either RandomGeneratorType or RandomStateType"
             )
 
+        # Fail early when size is incompatible with parameters
+        size_len = get_vector_length(size)
+        if size_len:
+            for i, (param, param_ndim_supp) in enumerate(
+                zip(dist_params, self.ndims_params)
+            ):
+                param_batched_dims = param.ndim - param_ndim_supp
+                if param_batched_dims > size_len:
+                    raise ValueError(
+                        f"Size length is incompatible with batched dimensions of parameter {i} {param}:\n"
+                        f"len(size) = {size_len}, len(batched dims {param}) = {param_batched_dims}. "
+                        f"Size length must be 0 or >= {param_batched_dims}"
+                    )
+
         shape = self._infer_shape(size, dist_params)
         _, static_shape = infer_static_shape(shape)
+
         dtype = self.dtype or dtype
 
         if dtype == "floatX":

tests/tensor/random/test_op.py

@@ -217,3 +217,17 @@ def test_random_maker_ops_no_seed():
     z = function(inputs=[], outputs=[default_rng()])()
     aes_res = z[0]
     assert isinstance(aes_res, np.random.Generator)
+
+
+def test_RandomVariable_incompatible_size():
+    rv_op = RandomVariable("normal", 0, [0, 0], config.floatX, inplace=True)
+    with pytest.raises(
+        ValueError, match="Size length is incompatible with batched dimensions"
+    ):
+        rv_op(np.zeros((1, 3)), 1, size=(3,))
+
+    rv_op = RandomVariable("dirichlet", 0, [1], config.floatX, inplace=True)
+    with pytest.raises(
+        ValueError, match="Size length is incompatible with batched dimensions"
+    ):
+        rv_op(np.zeros((2, 4, 3)), 1, size=(4,))