Remove size from Distribution signatures (#5788)

* Remove size from .dist() signature Closes #5754 * Align (Half)Flat signatures with superclass * Don't mention size in the docstring Co-authored-by: Ricardo Vieira <[email protected]> * Revert changes in `rng_fn` signatures Co-authored-by: Ricardo Vieira <[email protected]> Co-authored-by: Ricardo Vieira <[email protected]>
pymc-devs · May 23, 2022 · b5a5b56 · b5a5b56
1 parent d0af6b1
commit b5a5b56
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py
@@ -364,8 +364,8 @@ def __new__(cls, *args, **kwargs):
         return super().__new__(cls, *args, **kwargs)
 
     @classmethod
-    def dist(cls, *, size=None, **kwargs):
-        res = super().dist([], size=size, **kwargs)
+    def dist(cls, **kwargs):
+        res = super().dist([], **kwargs)
         return res
 
     def moment(rv, size):
@@ -432,8 +432,8 @@ def __new__(cls, *args, **kwargs):
         return super().__new__(cls, *args, **kwargs)
 
     @classmethod
-    def dist(cls, *, size=None, **kwargs):
-        res = super().dist([], size=size, **kwargs)
+    def dist(cls, **kwargs):
+        res = super().dist([], **kwargs)
         return res
 
     def moment(rv, size):

diff --git a/pymc/distributions/distribution.py b/pymc/distributions/distribution.py
@@ -226,8 +226,8 @@ def __new__(
         transform : optional
             See ``Model.register_rv``.
         **kwargs
-            Keyword arguments that will be forwarded to ``.dist()``.
-            Most prominently: ``shape`` and ``size``
+            Keyword arguments that will be forwarded to ``.dist()`` or the Aesara RV Op.
+            Most prominently: ``shape`` for ``.dist()`` or ``dtype`` for the Op.
 
         Returns
         -------
@@ -298,7 +298,6 @@ def dist(
         dist_params,
         *,
         shape: Optional[Shape] = None,
-        size: Optional[Size] = None,
         **kwargs,
     ) -> RandomVariable:
         """Creates a RandomVariable corresponding to the `cls` distribution.
@@ -312,8 +311,9 @@ def dist(
 
             An Ellipsis (...) may be inserted in the last position to short-hand refer to
             all the dimensions that the RV would get if no shape/size/dims were passed at all.
-        size : int, tuple, Variable, optional
-            For creating the RV like in Aesara/NumPy.
+        **kwargs
+            Keyword arguments that will be forwarded to the Aesara RV Op.
+            Most prominently: ``size`` or ``dtype``.
 
         Returns
         -------
@@ -337,6 +337,7 @@ def dist(
 
         if "dims" in kwargs:
             raise NotImplementedError("The use of a `.dist(dims=...)` API is not supported.")
+        size = kwargs.pop("size", None)
         if shape is not None and size is not None:
             raise ValueError(
                 f"Passing both `shape` ({shape}) and `size` ({size}) is not supported!"