Skip to content

Commit 80526ae

Browse files
Remove size from .dist() signature
Closes #5754
1 parent b91283e commit 80526ae

File tree

1 file changed

+6
-5
lines changed

1 file changed

+6
-5
lines changed

pymc/distributions/distribution.py

+6-5
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)