@@ -226,8 +226,8 @@ def __new__(
226226 transform : optional
227227 See ``Model.register_rv``.
228228 **kwargs
229- Keyword arguments that will be forwarded to ``.dist()``.
230- Most prominently: ``shape`` and ``size``
229+ Keyword arguments that will be forwarded to ``.dist()`` or the Aesara RV Op .
230+ Most prominently: ``shape`` for ``.dist()`` and ``size`` or ``dtype`` for the Op.
231231
232232 Returns
233233 -------
@@ -298,7 +298,6 @@ def dist(
298298 dist_params ,
299299 * ,
300300 shape : Optional [Shape ] = None ,
301- size : Optional [Size ] = None ,
302301 ** kwargs ,
303302 ) -> RandomVariable :
304303 """Creates a RandomVariable corresponding to the `cls` distribution.
@@ -312,8 +311,9 @@ def dist(
312311
313312 An Ellipsis (...) may be inserted in the last position to short-hand refer to
314313 all the dimensions that the RV would get if no shape/size/dims were passed at all.
315- size : int, tuple, Variable, optional
316- For creating the RV like in Aesara/NumPy.
314+ **kwargs
315+ Keyword arguments that will be forwarded to the Aesara RV Op.
316+ Most prominently: ``size`` or ``dtype``.
317317
318318 Returns
319319 -------
@@ -337,6 +337,7 @@ def dist(
337337
338338 if "dims" in kwargs :
339339 raise NotImplementedError ("The use of a `.dist(dims=...)` API is not supported." )
340+ size = kwargs .pop ("size" , None )
340341 if shape is not None and size is not None :
341342 raise ValueError (
342343 f"Passing both `shape` ({ shape } ) and `size` ({ size } ) is not supported!"
0 commit comments