pytensor 2.13.1__tar.gz → 2.14.0__tar.gz

{pytensor-2.13.1 → pytensor-2.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pytensor
-Version: 2.13.1
+Version: 2.14.0
 Summary: Optimizing compiler for evaluating mathematical expressions on CPUs and GPUs.
 Author-email: pymc-devs <pymc.devs@gmail.com>
 License: .. _license:
@@ -66,11 +66,10 @@ Classifier: Operating System :: POSIX
 Classifier: Operating System :: Unix
 Classifier: Operating System :: MacOS
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Python: >=3.8
+Requires-Python: >=3.9
 Description-Content-Type: text/x-rst
 Provides-Extra: complete
 Provides-Extra: development

{pytensor-2.13.1 → pytensor-2.14.0}/pyproject.toml

@@ -12,7 +12,7 @@ name = "pytensor"
 dynamic = [
   'version'
 ]
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 authors = [
     {name = "pymc-devs", email = "pymc.devs@gmail.com"}
 ]
@@ -34,7 +34,6 @@ classifiers = [
     "Operating System :: Unix",
     "Operating System :: MacOS",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.8",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
@@ -87,6 +86,7 @@ tests = [
     "pytest-cov>=2.6.1",
     "coverage>=5.1",
     "pytest-benchmark",
+    "pytest-mock",
 ]
 rtd = [
     "sphinx>=5.1.0,<6",

{pytensor-2.13.1 → pytensor-2.14.0}/pytensor/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2023-07-15T09:45:44+0200",
+ "date": "2023-07-26T10:12:05+0200",
  "dirty": false,
  "error": null,
- "full-revisionid": "9be43d07c50be420cceaa0434b48c15cfbda9121",
- "version": "2.13.1"
+ "full-revisionid": "14d2454c2a1f14c0fad6e39ec098031003809220",
+ "version": "2.14.0"
 }
 '''  # END VERSION_JSON
 

{pytensor-2.13.1 → pytensor-2.14.0}/pytensor/compile/function/pfunc.py

@@ -4,7 +4,7 @@ Provide a simple user friendly API.
 """
 
 from copy import copy
-from typing import Optional
+from typing import Optional, Sequence, Union, overload
 
 from pytensor.compile.function.types import Function, UnusedInputError, orig_function
 from pytensor.compile.io import In, Out
@@ -15,8 +15,9 @@ from pytensor.graph.basic import Constant, Variable, clone_node_and_cache
 from pytensor.graph.fg import FunctionGraph
 
 
+@overload
 def rebuild_collect_shared(
-    outputs,
+    outputs: Variable,
     inputs=None,
     replace=None,
     updates=None,
@@ -24,7 +25,107 @@ def rebuild_collect_shared(
     copy_inputs_over=True,
     no_default_updates=False,
     clone_inner_graphs=False,
-):
+) -> tuple[
+    list[Variable],
+    Variable,
+    tuple[
+        dict[Variable, Variable],
+        dict[SharedVariable, Variable],
+        list[Variable],
+        list[SharedVariable],
+    ],
+]:
+    ...
+
+
+@overload
+def rebuild_collect_shared(
+    outputs: Sequence[Variable],
+    inputs=None,
+    replace=None,
+    updates=None,
+    rebuild_strict=True,
+    copy_inputs_over=True,
+    no_default_updates=False,
+    clone_inner_graphs=False,
+) -> tuple[
+    list[Variable],
+    list[Variable],
+    tuple[
+        dict[Variable, Variable],
+        dict[SharedVariable, Variable],
+        list[Variable],
+        list[SharedVariable],
+    ],
+]:
+    ...
+
+
+@overload
+def rebuild_collect_shared(
+    outputs: Out,
+    inputs=None,
+    replace=None,
+    updates=None,
+    rebuild_strict=True,
+    copy_inputs_over=True,
+    no_default_updates=False,
+    clone_inner_graphs=False,
+) -> tuple[
+    list[Variable],
+    Out,
+    tuple[
+        dict[Variable, Variable],
+        dict[SharedVariable, Variable],
+        list[Variable],
+        list[SharedVariable],
+    ],
+]:
+    ...
+
+
+@overload
+def rebuild_collect_shared(
+    outputs: Sequence[Out],
+    inputs=None,
+    replace=None,
+    updates=None,
+    rebuild_strict=True,
+    copy_inputs_over=True,
+    no_default_updates=False,
+    clone_inner_graphs=False,
+) -> tuple[
+    list[Variable],
+    list[Out],
+    tuple[
+        dict[Variable, Variable],
+        dict[SharedVariable, Variable],
+        list[Variable],
+        list[SharedVariable],
+    ],
+]:
+    ...
+
+
+def rebuild_collect_shared(
+    outputs: Union[Sequence[Variable], Variable, Out, Sequence[Out]],
+    inputs=None,
+    replace=None,
+    updates=None,
+    rebuild_strict=True,
+    copy_inputs_over=True,
+    no_default_updates=False,
+    clone_inner_graphs=False,
+) -> tuple[
+    list[Variable],
+    Union[list[Variable], Variable, Out, list[Out]],
+    tuple[
+        dict[Variable, Variable],
+        dict[SharedVariable, Variable],
+        list[Variable],
+        list[SharedVariable],
+    ],
+]:
     r"""Replace subgraphs of a computational graph.
 
     It returns a set of dictionaries and lists which collect (partial?)
@@ -260,7 +361,7 @@ def rebuild_collect_shared(
     return (
         input_variables,
         cloned_outputs,
-        [clone_d, update_d, update_expr, shared_inputs],
+        (clone_d, update_d, update_expr, shared_inputs),
     )
 
 

{pytensor-2.13.1 → pytensor-2.14.0}/pytensor/graph/basic.py

@@ -1003,14 +1003,14 @@ def applys_between(
 def truncated_graph_inputs(
     outputs: Sequence[Variable],
     ancestors_to_include: Optional[Collection[Variable]] = None,
-) -> List[Variable]:
+) -> list[Variable]:
     """Get the truncate graph inputs.
 
     Unlike :func:`graph_inputs` this function will return
-    the closest nodes to outputs that do not depend on
+    the closest variables to outputs that do not depend on
     ``ancestors_to_include``. So given all the returned
-    variables provided there is no missing node to
-    compute the output and all nodes are independent
+    variables provided there is no missing variable to
+    compute the output and all variables are independent
     from each other.
 
     Parameters
@@ -1027,7 +1027,7 @@ def truncated_graph_inputs(
 
     Examples
     --------
-    The returned nodes marked in (parenthesis), ancestors nodes are ``c``, output nodes are ``o``
+    The returned variables marked in (parenthesis), ancestors variables are ``c``, output variables are ``o``
 
     * No ancestors to include
 
@@ -1047,7 +1047,7 @@ def truncated_graph_inputs(
 
         (c) - (c) - o
 
-    * Additional nodes are present
+    * Additional variables are present
 
     .. code-block::
 
@@ -1076,58 +1076,60 @@ def truncated_graph_inputs(
 
     """
     # simple case, no additional ancestors to include
-    truncated_inputs = list()
-    # blockers have known independent nodes and ancestors to include
+    truncated_inputs: list[Variable] = list()
+    # blockers have known independent variables and ancestors to include
     candidates = list(outputs)
     if not ancestors_to_include:  # None or empty
         # just filter out unique variables
-        for node in candidates:
-            if node not in truncated_inputs:
-                truncated_inputs.append(node)
+        for variable in candidates:
+            if variable not in truncated_inputs:
+                truncated_inputs.append(variable)
         # no more actions are needed
         return truncated_inputs
 
-    blockers: Set[Variable] = set(ancestors_to_include)
-    # enforce O(1) check for node in ancestors to include
+    blockers: set[Variable] = set(ancestors_to_include)
+    # variables that go here are under check already, do not repeat the loop for them
+    seen: set[Variable] = set()
+    # enforce O(1) check for variable in ancestors to include
     ancestors_to_include = blockers.copy()
 
     while candidates:
         # on any new candidate
-        node = candidates.pop()
-
-        # There was a repeated reference to this node, we have already investigated it
-        if node in truncated_inputs:
+        variable = candidates.pop()
+        # we've looked into this variable already
+        if variable in seen:
             continue
-
-        # check if the node is independent, never go above blockers;
-        # blockers are independent nodes and ancestors to include
-        if node in ancestors_to_include:
-            # The case where node is in ancestors to include so we check if it depends on others
+        # check if the variable is independent, never go above blockers;
+        # blockers are independent variables and ancestors to include
+        elif variable in ancestors_to_include:
+            # The case where variable is in ancestors to include so we check if it depends on others
             # it should be removed from the blockers to check against the rest
-            dependent = variable_depends_on(node, ancestors_to_include - {node})
+            dependent = variable_depends_on(variable, ancestors_to_include - {variable})
             # ancestors to include that are present in the graph (not disconnected)
             # should be added to truncated_inputs
-            truncated_inputs.append(node)
+            truncated_inputs.append(variable)
             if dependent:
                 # if the ancestors to include is still dependent we need to go above, the search is not yet finished
-                # owner can never be None for a dependent node
-                candidates.extend(node.owner.inputs)
+                # owner can never be None for a dependent variable
+                candidates.extend(n for n in variable.owner.inputs if n not in seen)
         else:
-            # A regular node to check
-            dependent = variable_depends_on(node, blockers)
-            # all regular nodes fall to blockers
+            # A regular variable to check
+            dependent = variable_depends_on(variable, blockers)
+            # all regular variables fall to blockers
             # 1. it is dependent - further search irrelevant
-            # 2. it is independent - the search node is inside the closure
-            blockers.add(node)
-            # if we've found an independent node and it is not in blockers so far
-            # it is a new independent node not present in ancestors to include
+            # 2. it is independent - the search variable is inside the closure
+            blockers.add(variable)
+            # if we've found an independent variable and it is not in blockers so far
+            # it is a new independent variable not present in ancestors to include
             if dependent:
-                # populate search if it's not an independent node
-                # owner can never be None for a dependent node
-                candidates.extend(node.owner.inputs)
+                # populate search if it's not an independent variable
+                # owner can never be None for a dependent variable
+                candidates.extend(n for n in variable.owner.inputs if n not in seen)
             else:
                 # otherwise, do not search beyond
-                truncated_inputs.append(node)
+                truncated_inputs.append(variable)
+        # add variable to seen, no point in checking it once more
+        seen.add(variable)
     return truncated_inputs
 
 

{pytensor-2.13.1 → pytensor-2.14.0}/pytensor/graph/replace.py

@@ -1,27 +1,56 @@
 from functools import partial
-from typing import (
-    Collection,
-    Dict,
-    Iterable,
-    List,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-    cast,
-)
-
-from pytensor.graph.basic import Constant, Variable, truncated_graph_inputs
+from typing import Iterable, Optional, Sequence, Union, cast, overload
+
+from pytensor.graph.basic import Apply, Constant, Variable, truncated_graph_inputs
 from pytensor.graph.fg import FunctionGraph
 
 
+ReplaceTypes = Union[Iterable[tuple[Variable, Variable]], dict[Variable, Variable]]
+
+
+def _format_replace(replace: Optional[ReplaceTypes] = None) -> dict[Variable, Variable]:
+    items: dict[Variable, Variable]
+    if isinstance(replace, dict):
+        # PyLance has issues with type resolution
+        items = cast(dict[Variable, Variable], replace)
+    elif isinstance(replace, Iterable):
+        items = dict(replace)
+    elif replace is None:
+        items = {}
+    else:
+        raise ValueError(
+            "replace is neither a dictionary, list, "
+            f"tuple or None ! The value provided is {replace},"
+            f"of type {type(replace)}"
+        )
+    return items
+
+
+@overload
+def clone_replace(
+    output: Sequence[Variable],
+    replace: Optional[ReplaceTypes] = None,
+    **rebuild_kwds,
+) -> list[Variable]:
+    ...
+
+
+@overload
 def clone_replace(
-    output: Collection[Variable],
+    output: Variable,
     replace: Optional[
-        Union[Iterable[Tuple[Variable, Variable]], Dict[Variable, Variable]]
+        Union[Iterable[tuple[Variable, Variable]], dict[Variable, Variable]]
     ] = None,
     **rebuild_kwds,
-) -> List[Variable]:
+) -> Variable:
+    ...
+
+
+def clone_replace(
+    output: Union[Sequence[Variable], Variable],
+    replace: Optional[ReplaceTypes] = None,
+    **rebuild_kwds,
+) -> Union[list[Variable], Variable]:
     """Clone a graph and replace subgraphs within it.
 
     It returns a copy of the initial subgraph with the corresponding
@@ -39,19 +68,8 @@ def clone_replace(
     """
     from pytensor.compile.function.pfunc import rebuild_collect_shared
 
-    items: Union[List[Tuple[Variable, Variable]], Tuple[Tuple[Variable, Variable], ...]]
-    if isinstance(replace, dict):
-        items = list(replace.items())
-    elif isinstance(replace, (list, tuple)):
-        items = replace
-    elif replace is None:
-        items = []
-    else:
-        raise ValueError(
-            "replace is neither a dictionary, list, "
-            f"tuple or None ! The value provided is {replace},"
-            f"of type {type(replace)}"
-        )
+    items = list(_format_replace(replace).items())
+
     tmp_replace = [(x, x.type()) for x, y in items]
     new_replace = [(x, y) for ((_, x), (_, y)) in zip(tmp_replace, items)]
     _, _outs, _ = rebuild_collect_shared(output, [], tmp_replace, [], **rebuild_kwds)
@@ -59,20 +77,40 @@ def clone_replace(
     # TODO Explain why we call it twice ?!
     _, outs, _ = rebuild_collect_shared(_outs, [], new_replace, [], **rebuild_kwds)
 
-    return cast(List[Variable], outs)
+    return outs
 
 
+@overload
+def graph_replace(
+    outputs: Variable,
+    replace: Optional[ReplaceTypes] = None,
+    *,
+    strict=True,
+) -> Variable:
+    ...
+
+
+@overload
 def graph_replace(
     outputs: Sequence[Variable],
-    replace: Dict[Variable, Variable],
+    replace: Optional[ReplaceTypes] = None,
+    *,
+    strict=True,
+) -> list[Variable]:
+    ...
+
+
+def graph_replace(
+    outputs: Union[Sequence[Variable], Variable],
+    replace: Optional[ReplaceTypes] = None,
     *,
     strict=True,
-) -> List[Variable]:
+) -> Union[list[Variable], Variable]:
     """Replace variables in ``outputs`` by ``replace``.
 
     Parameters
     ----------
-    outputs: Sequence[Variable]
+    outputs: Union[Sequence[Variable], Variable]
         Output graph
     replace: Dict[Variable, Variable]
         Replace mapping
@@ -83,20 +121,26 @@ def graph_replace(
 
     Returns
     -------
-    List[Variable]
-        Output graph with subgraphs replaced
+    Union[Variable, List[Variable]]
+        Output graph with subgraphs replaced, see function overload for the exact type
 
     Raises
     ------
     ValueError
-        If some replacemens could not be applied and strict is True
+        If some replacements could not be applied and strict is True
     """
+    as_list = False
+    if not isinstance(outputs, Sequence):
+        outputs = [outputs]
+    else:
+        as_list = True
+    replace_dict = _format_replace(replace)
     # collect minimum graph inputs which is required to compute outputs
     # and depend on replacements
     # additionally remove constants, they do not matter in clone get equiv
     conditions = [
         c
-        for c in truncated_graph_inputs(outputs, replace)
+        for c in truncated_graph_inputs(outputs, replace_dict)
         if not isinstance(c, Constant)
     ]
     # for the function graph we need the clean graph where
@@ -117,7 +161,7 @@ def graph_replace(
     # replace the conditions back
     fg_replace = {equiv[c]: c for c in conditions}
     # add the replacements on top of input mappings
-    fg_replace.update({equiv[r]: v for r, v in replace.items() if r in equiv})
+    fg_replace.update({equiv[r]: v for r, v in replace_dict.items() if r in equiv})
     # replacements have to be done in reverse topological order so that nested
     # expressions get recursively replaced correctly
 
@@ -126,12 +170,14 @@ def graph_replace(
     # So far FunctionGraph does these replacements inplace it is thus unsafe
     # apply them using fg.replace, it may change the original graph
     if strict:
-        non_fg_replace = {r: v for r, v in replace.items() if r not in equiv}
+        non_fg_replace = {r: v for r, v in replace_dict.items() if r not in equiv}
         if non_fg_replace:
             raise ValueError(f"Some replacements were not used: {non_fg_replace}")
     toposort = fg.toposort()
 
-    def toposort_key(fg: FunctionGraph, ts, pair):
+    def toposort_key(
+        fg: FunctionGraph, ts: list[Apply], pair: tuple[Variable, Variable]
+    ) -> int:
         key, _ = pair
         if key.owner is not None:
             return ts.index(key.owner)
@@ -148,4 +194,7 @@ def graph_replace(
         reverse=True,
     )
     fg.replace_all(sorted_replacements, import_missing=True)
-    return list(fg.outputs)
+    if as_list:
+        return list(fg.outputs)
+    else:
+        return fg.outputs[0]

{pytensor-2.13.1 → pytensor-2.14.0}/pytensor/tensor/random/basic.py

@@ -5,10 +5,13 @@ import numpy as np
 import scipy.stats as stats
 
 import pytensor
-from pytensor.tensor.basic import as_tensor_variable
-from pytensor.tensor.random.op import RandomVariable, default_supp_shape_from_params
+from pytensor.tensor.basic import arange, as_tensor_variable
+from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.random.type import RandomGeneratorType, RandomStateType
-from pytensor.tensor.random.utils import broadcast_params
+from pytensor.tensor.random.utils import (
+    broadcast_params,
+    supp_shape_from_ref_param_shape,
+)
 from pytensor.tensor.random.var import (
     RandomGeneratorSharedVariable,
     RandomStateSharedVariable,
@@ -855,6 +858,14 @@ class MvNormalRV(RandomVariable):
     dtype = "floatX"
     _print_name = ("MultivariateNormal", "\\operatorname{MultivariateNormal}")
 
+    def _supp_shape_from_params(self, dist_params, param_shapes=None):
+        return supp_shape_from_ref_param_shape(
+            ndim_supp=self.ndim_supp,
+            dist_params=dist_params,
+            param_shapes=param_shapes,
+            ref_param_idx=0,
+        )
+
     def __call__(self, mean=None, cov=None, size=None, **kwargs):
         r""" "Draw samples from a multivariate normal distribution.
 
@@ -933,6 +944,14 @@ class DirichletRV(RandomVariable):
     dtype = "floatX"
     _print_name = ("Dirichlet", "\\operatorname{Dirichlet}")
 
+    def _supp_shape_from_params(self, dist_params, param_shapes=None):
+        return supp_shape_from_ref_param_shape(
+            ndim_supp=self.ndim_supp,
+            dist_params=dist_params,
+            param_shapes=param_shapes,
+            ref_param_idx=0,
+        )
+
     def __call__(self, alphas, size=None, **kwargs):
         r"""Draw samples from a dirichlet distribution.
 
@@ -1776,9 +1795,12 @@ class MultinomialRV(RandomVariable):
         """
         return super().__call__(n, p, size=size, **kwargs)
 
-    def _supp_shape_from_params(self, dist_params, rep_param_idx=1, param_shapes=None):
-        return default_supp_shape_from_params(
-            self.ndim_supp, dist_params, rep_param_idx, param_shapes
+    def _supp_shape_from_params(self, dist_params, param_shapes=None):
+        return supp_shape_from_ref_param_shape(
+            ndim_supp=self.ndim_supp,
+            dist_params=dist_params,
+            param_shapes=param_shapes,
+            ref_param_idx=1,
         )
 
     @classmethod
@@ -2050,18 +2072,15 @@ class PermutationRV(RandomVariable):
 
     @classmethod
     def rng_fn(cls, rng, x, size):
-        return rng.permutation(x if x.ndim > 0 else x.item())
-
-    def _infer_shape(self, size, dist_params, param_shapes=None):
-        param_shapes = param_shapes or [p.shape for p in dist_params]
-
-        (x,) = dist_params
-        (x_shape,) = param_shapes
-
-        if x.ndim == 0:
-            return (x,)
-        else:
-            return x_shape
+        return rng.permutation(x)
+
+    def _supp_shape_from_params(self, dist_params, param_shapes=None):
+        return supp_shape_from_ref_param_shape(
+            ndim_supp=self.ndim_supp,
+            dist_params=dist_params,
+            param_shapes=param_shapes,
+            ref_param_idx=0,
+        )
 
     def __call__(self, x, **kwargs):
         r"""Randomly permute a sequence or a range of values.
@@ -2074,15 +2093,35 @@ class PermutationRV(RandomVariable):
         Parameters
         ----------
         x
-            If `x` is an integer, randomly permute `np.arange(x)`. If `x` is a sequence,
-            shuffle its elements randomly.
+            Elements to be shuffled.
 
         """
         x = as_tensor_variable(x)
         return super().__call__(x, dtype=x.dtype, **kwargs)
 
 
-permutation = PermutationRV()
+_permutation = PermutationRV()
+
+
+def permutation(x, **kwargs):
+    r"""Randomly permute a sequence or a range of values.
+
+    Signature
+    ---------
+
+    `(x) -> (x)`
+
+    Parameters
+    ----------
+    x
+        If `x` is an integer, randomly permute `np.arange(x)`. If `x` is a sequence,
+        shuffle its elements randomly.
+
+    """
+    x = as_tensor_variable(x)
+    if x.type.ndim == 0:
+        x = arange(x)
+    return _permutation(x, **kwargs)
 
 
 __all__ = [