Extract RandomVariable size parameter handling to a separate utility function

aesara/tensor/random/op.py

@@ -11,15 +11,14 @@ from aesara.graph.op import Op
 from aesara.misc.safe_asarray import _asarray
 from aesara.tensor.basic import (
     as_tensor_variable,
-    cast,
     constant,
     get_scalar_constant_value,
     get_vector_length,
 )
 from aesara.tensor.exceptions import NotScalarConstantError
 from aesara.tensor.random.type import RandomStateType
-from aesara.tensor.random.utils import params_broadcast_shapes
-from aesara.tensor.type import TensorType, all_dtypes, int_dtypes
+from aesara.tensor.random.utils import normalize_size_param, params_broadcast_shapes
+from aesara.tensor.type import TensorType, all_dtypes
 from aesara.tensor.type_other import NoneConst
 
 
@@ -348,18 +347,7 @@ class RandomVariable(Op):
             `(rng_var, out_var)`.
 
         """
-        if size is None:
-            size = constant([], dtype="int64")
-        elif isinstance(size, int):
-            size = as_tensor_variable([size], ndim=1)
-        elif not isinstance(size, (np.ndarray, Variable, Sequence)):
-            raise TypeError(
-                "Parameter size must be None, an integer, or a sequence with integers."
-            )
-        else:
-            size = cast(as_tensor_variable(size, ndim=1), "int64")
-
-        assert size.dtype in int_dtypes
+        size = normalize_size_param(size)
 
         dist_params = tuple(
             as_tensor_variable(p) if not isinstance(p, Variable) else p

aesara/tensor/random/utils.py

+from collections.abc import Sequence
 from functools import wraps
 from itertools import zip_longest
 
@@ -5,8 +6,10 @@ import numpy as np
 
 from aesara.compile.sharedvalue import shared
 from aesara.graph.basic import Variable
+from aesara.tensor.basic import as_tensor_variable, cast, constant
 from aesara.tensor.extra_ops import broadcast_to
 from aesara.tensor.math import maximum
+from aesara.tensor.type import int_dtypes
 
 
 def params_broadcast_shapes(param_shapes, ndims_params, use_aesara=True):
@@ -95,6 +98,24 @@ def broadcast_params(params, ndims_params):
     return bcast_params
 
 
+def normalize_size_param(size):
+    """Create an Aesara value for a ``RandomVariable`` ``size`` parameter."""
+    if size is None:
+        size = constant([], dtype="int64")
+    elif isinstance(size, int):
+        size = as_tensor_variable([size], ndim=1)
+    elif not isinstance(size, (np.ndarray, Variable, Sequence)):
+        raise TypeError(
+            "Parameter size must be None, an integer, or a sequence with integers."
+        )
+    else:
+        size = cast(as_tensor_variable(size, ndim=1), "int64")
+
+    assert size.dtype in int_dtypes
+
+    return size
+
+
 class RandomStream:
     """Module component with similar interface to `numpy.random.RandomState`.