From 9978e9e23a46b7afa2df6e984def5fcc26c8bef7 Mon Sep 17 00:00:00 2001 From: UditNayak Date: Sat, 11 Nov 2023 04:37:32 +0530 Subject: [PATCH 1/3] Add missing functions in shape_utils doc, improve return formatting, and include type hints in docstring #7004 --- docs/source/api/shape_utils.rst | 1 + pymc/distributions/shape_utils.py | 78 ++++++++++++++++++++----------- 2 files changed, 53 insertions(+), 26 deletions(-) diff --git a/docs/source/api/shape_utils.rst b/docs/source/api/shape_utils.rst index 7f78052f87..28e49e78f6 100644 --- a/docs/source/api/shape_utils.rst +++ b/docs/source/api/shape_utils.rst @@ -16,3 +16,4 @@ This module introduces functions that are made aware of the requested `size_tupl to_tuple rv_size_is_none change_dist_size + broadcast_dist_samples_shape diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index f85e5dc477..38cd05077f 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -20,7 +20,7 @@ import warnings from functools import singledispatch -from typing import Any, Optional, Sequence, Tuple, Union, cast +from typing import Any, Iterable, Optional, Sequence, Tuple, Union, cast import numpy as np @@ -49,18 +49,28 @@ from pymc.util import _add_future_warning_tag -def to_tuple(shape): - """Convert ints, arrays, and Nones to tuples +def to_tuple(shape: Optional[Union[None, int, np.ndarray]]) -> Tuple: + """Convert integers, arrays, and Nones to tuples. Parameters ---------- - shape: None, int or array-like - Represents the shape to convert to tuple. + shape : None, int, or array-like + Represents the shape to convert to a tuple. Returns ------- - If `shape` is None, returns an empty tuple. If it's an int, (shape,) is - returned. If it is array-like, tuple(shape) is returned. + tuple + If `shape` is None, returns an empty tuple. If it's an int, (shape,) is + returned. If it is array-like, `tuple(shape)` is returned. + + Examples + -------- + >>> to_tuple(None) + () + >>> to_tuple(5) + (5,) + >>> to_tuple([1, 2, 3]) + (1, 2, 3) """ if shape is None: return tuple() @@ -87,26 +97,25 @@ def _check_shape_type(shape): return tuple(out) -def broadcast_dist_samples_shape(shapes, size=None): - """Apply shape broadcasting to shape tuples but assuming that the shapes - correspond to draws from random variables, with the `size` tuple possibly - prepended to it. The `size` prepend is ignored to consider if the supplied - `shapes` can broadcast or not. It is prepended to the resulting broadcasted - `shapes`, if any of the shape tuples had the `size` prepend. +def broadcast_dist_samples_shape(shapes: Iterable[Tuple[int, ...]], size: Optional[int] = None) -> Tuple[int, ...]: + """Apply shape broadcasting to shape tuples for random variables. Parameters ---------- - shapes: Iterable of tuples holding the distribution samples shapes - size: None, int or tuple (optional) - size of the sample set requested. + shapes : Iterable of tuples + Tuples holding the distribution samples shapes. + size : None, int, or tuple, optional + Size of the sample set requested. Returns ------- - tuple of the resulting shape + tuple + The resulting broadcasted shape. Examples -------- .. code-block:: python + size = 100 shape0 = (size,) shape1 = (size, 5) @@ -114,7 +123,9 @@ def broadcast_dist_samples_shape(shapes, size=None): out = broadcast_dist_samples_shape([shape0, shape1, shape2], size=size) assert out == (size, 4, 5) + .. code-block:: python + size = 100 shape0 = (size,) shape1 = (5,) @@ -122,7 +133,9 @@ def broadcast_dist_samples_shape(shapes, size=None): out = broadcast_dist_samples_shape([shape0, shape1, shape2], size=size) assert out == (size, 4, 5) + .. code-block:: python + size = 100 shape0 = (1,) shape1 = (5,) @@ -291,7 +304,18 @@ def find_size( def rv_size_is_none(size: Variable) -> bool: - """Check whether an rv size is None (ie., pt.Constant([]))""" + """Check whether the size of a random variable is None. + + Parameters + ---------- + size : Variable + The size variable to check. + + Returns + ------- + bool + True if the size is None (i.e., pt.Constant([])), False otherwise. + """ return size.type.shape == (0,) # type: ignore [attr-defined] @@ -311,19 +335,21 @@ def change_dist_size( Parameters ---------- - dist: + dist : TensorVariable The old distribution to be resized. - new_size: + new_size : Union[int, Tuple[int, ...]] The new size of the distribution. - expand: bool, optional - If True, `new_size` is prepended to the existing distribution `size`, so that - the final size is equal to (*new_size, *dist.size). Defaults to false. + expand : bool, optional + If True, `new_size` is prepended to the existing distribution `size`, + so that the final size is equal to (*new_size, *dist.size). + Defaults to False. Returns ------- - A new distribution variable that is equivalent to the original distribution with - the new size. The new distribution will not reuse the old RandomState/Generator - input, so it will be independent from the original distribution. + TensorVariable + A new distribution variable equivalent to the original distribution + with the new size. The new distribution will not reuse the old + RandomState/Generator input, making it independent from the original. Examples -------- From e84acd887ae7767dc7df8badf52ed0536e11fc9f Mon Sep 17 00:00:00 2001 From: UditNayak Date: Sat, 11 Nov 2023 11:28:19 +0530 Subject: [PATCH 2/3] fix pre-commit --- pymc/distributions/shape_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index 38cd05077f..c63036f244 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -62,7 +62,7 @@ def to_tuple(shape: Optional[Union[None, int, np.ndarray]]) -> Tuple: tuple If `shape` is None, returns an empty tuple. If it's an int, (shape,) is returned. If it is array-like, `tuple(shape)` is returned. - + Examples -------- >>> to_tuple(None) @@ -97,7 +97,9 @@ def _check_shape_type(shape): return tuple(out) -def broadcast_dist_samples_shape(shapes: Iterable[Tuple[int, ...]], size: Optional[int] = None) -> Tuple[int, ...]: +def broadcast_dist_samples_shape( + shapes: Iterable[Tuple[int, ...]], size: Optional[int] = None +) -> Tuple[int, ...]: """Apply shape broadcasting to shape tuples for random variables. Parameters From 464ac0a1b8675184917a8fc5877aa60883c115fb Mon Sep 17 00:00:00 2001 From: UditNayak Date: Sun, 12 Nov 2023 00:14:58 +0530 Subject: [PATCH 3/3] Add TensorVariable & remove broadcast_dist_samples_shape --- docs/source/api/shape_utils.rst | 1 - pymc/distributions/shape_utils.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/source/api/shape_utils.rst b/docs/source/api/shape_utils.rst index 28e49e78f6..7f78052f87 100644 --- a/docs/source/api/shape_utils.rst +++ b/docs/source/api/shape_utils.rst @@ -16,4 +16,3 @@ This module introduces functions that are made aware of the requested `size_tupl to_tuple rv_size_is_none change_dist_size - broadcast_dist_samples_shape diff --git a/pymc/distributions/shape_utils.py b/pymc/distributions/shape_utils.py index c63036f244..74f7ebbb6f 100644 --- a/pymc/distributions/shape_utils.py +++ b/pymc/distributions/shape_utils.py @@ -339,7 +339,7 @@ def change_dist_size( ---------- dist : TensorVariable The old distribution to be resized. - new_size : Union[int, Tuple[int, ...]] + new_size : Union[int, Tuple[int, ...], TensorVariable] The new size of the distribution. expand : bool, optional If True, `new_size` is prepended to the existing distribution `size`,