Merge branch 'main' into pylint

Author: michaelosthege

.github/release.yml

@@ -0,0 +1,26 @@
+# This file has been mostly taken verbatim from https://github.com/pymc-devs/pymc/blob/main/.github/release.yml
+#
+# This file contains configuration for the automatic generation of release notes in GitHub.
+# It's not perfect, but it makes it a little less laborious to write informative release notes.
+# Also see https://docs.github.com/en/repositories/releasing-projects-on-github/automatically-generated-release-notes
+changelog:
+  exclude:
+    labels:
+      - no releasenotes
+  categories:
+    - title: Major Changes 🛠
+      labels:
+        - major
+    - title: New Features & Bugfixes 🎉
+      labels:
+        - bug
+        - enhancements
+        - feature-request
+    - title: Docs & Maintenance 🔧
+      labels:
+        - docs
+        - installation
+        - maintenance
+        - pre-commit
+        - tests
+        - "*"

README.rst

@@ -114,10 +114,10 @@ Contributing
 We welcome bug reports and fixes and improvements to the documentation.
 
 For more information on contributing, please see the
-`contributing guide <https://github.com/pymc-devs/pytensor/CONTRIBUTING.md>`.
+`contributing guide <https://github.com/pymc-devs/pytensor/CONTRIBUTING.md>`__.
 
 A good place to start contributing is by looking through the issues
-`here <https://github.com/pymc-devs/pytensor/issues`.
+`here <https://github.com/pymc-devs/pytensor/issues>`__.
 
 
 .. |Project Name| replace:: PyTensor

doc/library/compile/shared.rst

@@ -13,7 +13,7 @@
 
 .. class:: SharedVariable
 
-    Variable with Storage that is shared between functions that it appears in.
+    Variable with storage that is shared between the compiled functions that it appears in.
     These variables are meant to be created by registered *shared constructors*
     (see :func:`shared_constructor`).
 
@@ -68,18 +68,17 @@
 
         A container to use for this SharedVariable when it is an implicit function parameter.
 
-        :type: class:`Container`
 
 .. autofunction:: shared
 
 .. function:: shared_constructor(ctor)
 
     Append `ctor` to the list of shared constructors (see :func:`shared`).
 
-    Each registered constructor ``ctor`` will be called like this:
+    Each registered constructor `ctor` will be called like this:
 
     .. code-block:: python
 
         ctor(value, name=name, strict=strict, **kwargs)
 
-    If it do not support given value, it must raise a TypeError.
+    If it do not support given value, it must raise a `TypeError`.

pytensor/compile/builders.py

@@ -2,7 +2,7 @@
 from collections import OrderedDict
 from copy import copy
 from functools import partial
-from typing import List, Optional, Sequence, cast
+from typing import Dict, List, Optional, Sequence, Tuple, cast
 
 import pytensor.tensor as at
 from pytensor import function
@@ -19,7 +19,6 @@
     clone_replace,
     graph_inputs,
     io_connection_pattern,
-    replace_nominals_with_dummies,
 )
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.null_type import NullType
@@ -82,6 +81,81 @@ def local_traverse(out):
     return ret
 
 
+def construct_nominal_fgraph(
+    inputs: Sequence[Variable], outputs: Sequence[Variable]
+) -> Tuple[
+    FunctionGraph,
+    Sequence[Variable],
+    Dict[Variable, Variable],
+    Dict[Variable, Variable],
+]:
+    """Construct an inner-`FunctionGraph` with ordered nominal inputs."""
+    dummy_inputs = []
+    for n, inp in enumerate(inputs):
+        if (
+            not isinstance(inp, Variable)
+            or isinstance(inp, Constant)
+            or isinstance(inp, SharedVariable)
+        ):
+            raise TypeError(
+                f"Inputs and outputs must be non-Constant/shared Variable instances; got {inp}"
+            )
+
+        dummy_inputs.append(inp.type())
+
+    dummy_shared_inputs = []
+    shared_inputs = []
+    for var in graph_inputs(outputs, inputs):
+        if isinstance(var, SharedVariable):
+            # To correctly support shared variables the inner-graph should
+            # not see them; otherwise, there will be problems with
+            # gradients.
+            # That's why we collect the shared variables and replace them
+            # with dummies.
+            shared_inputs.append(var)
+            dummy_shared_inputs.append(var.type())
+        elif var not in inputs and not isinstance(var, Constant):
+            raise MissingInputError(f"OpFromGraph is missing an input: {var}")
+
+    replacements = dict(zip(inputs + shared_inputs, dummy_inputs + dummy_shared_inputs))
+
+    new = rebuild_collect_shared(
+        cast(Sequence[Variable], outputs),
+        inputs=inputs + shared_inputs,
+        replace=replacements,
+        copy_inputs_over=False,
+    )
+    (
+        local_inputs,
+        local_outputs,
+        (clone_d, update_d, update_expr, new_shared_inputs),
+    ) = new
+
+    assert len(local_inputs) == len(inputs) + len(shared_inputs)
+    assert len(local_outputs) == len(outputs)
+    assert not update_d
+    assert not update_expr
+    assert not new_shared_inputs
+
+    fgraph = FunctionGraph(local_inputs, local_outputs, clone=False)
+
+    # The inputs need to be `NominalVariable`s so that we can merge
+    # inner-graphs
+    nominal_local_inputs = tuple(
+        NominalVariable(n, var.type) for n, var in enumerate(local_inputs)
+    )
+
+    fgraph.replace_all(zip(local_inputs, nominal_local_inputs))
+
+    for i, inp in enumerate(fgraph.inputs):
+        nom_inp = nominal_local_inputs[i]
+        fgraph.inputs[i] = nom_inp
+        fgraph.clients.pop(inp, None)
+        fgraph.add_input(nom_inp)
+
+    return fgraph, shared_inputs, update_d, update_expr
+
+
 class OpFromGraph(Op, HasInnerGraph):
     r"""
     This creates an `Op` from inputs and outputs lists of variables.
@@ -333,66 +407,21 @@ def __init__(
         if not (isinstance(inputs, list) and isinstance(outputs, list)):
             raise TypeError("Inputs and outputs must be lists")
 
-        for i in inputs + outputs:
-            if not isinstance(i, Variable):
+        for out in outputs:
+            if not isinstance(out, Variable):
                 raise TypeError(
-                    f"Inputs and outputs must be Variable instances; got {i}"
+                    f"Inputs and outputs must be Variable instances; got {out}"
                 )
-            if i in inputs:
-                if isinstance(i, Constant):
-                    raise TypeError(f"Constants not allowed as inputs; {i}")
-                if isinstance(i, SharedVariable):
-                    raise TypeError(f"SharedVariables not allowed as inputs; {i}")
-
-        for var in graph_inputs(outputs, inputs):
-            if var not in inputs and not isinstance(var, (Constant, SharedVariable)):
-                raise MissingInputError(f"OpFromGraph is missing an input: {var}")
 
         if "updates" in kwargs or "givens" in kwargs:
-            raise NotImplementedError("Updates and givens are not allowed here")
+            raise NotImplementedError("Updates and givens are not supported")
 
         self.is_inline = inline
 
-        # To correctly support shared variables the inner fct should
-        # not see them. Otherwise there is a problem with the gradient.
-        self.shared_inputs = []
-        for var in graph_inputs(outputs):
-            if isinstance(var, SharedVariable):
-                self.shared_inputs.append(var)
-
-        inputs, outputs = replace_nominals_with_dummies(inputs, outputs)
-
-        # The inputs should be `NominalVariable`s, so that graphs can be merged
-        replacements = {}
-        for n, v in enumerate(inputs):
-            replacements[v] = NominalVariable(n, v.type)
-
-        shared_vars = [
-            NominalVariable(n, var.type)
-            for n, var in enumerate(self.shared_inputs, start=len(inputs) + 1)
-        ]
-
-        replacements.update(dict(zip(self.shared_inputs, shared_vars)))
-
-        new = rebuild_collect_shared(
-            cast(Sequence[Variable], outputs),
-            inputs=inputs + shared_vars,
-            replace=replacements,
-            copy_inputs_over=False,
+        self.fgraph, self.shared_inputs, _, _ = construct_nominal_fgraph(
+            inputs, outputs
         )
-        (
-            local_inputs,
-            local_outputs,
-            (clone_d, update_d, update_expr, shared_inputs),
-        ) = new
-
-        assert len(local_inputs) == len(inputs) + len(self.shared_inputs)
-        assert len(local_outputs) == len(outputs)
-        assert not update_d
-        assert not update_expr
-        assert not shared_inputs
-
-        self.fgraph = FunctionGraph(local_inputs, local_outputs, clone=False)
+
         self.kwargs = kwargs
         self.input_types = [inp.type for inp in inputs]
         self.output_types = [out.type for out in outputs]
@@ -415,6 +444,7 @@ def __init__(
         else:
             self.set_lop_overrides("default")
             self._lop_type = "lop"
+
         self.set_rop_overrides(rop_overrides)
 
         self._connection_pattern = connection_pattern

pytensor/compile/function/pfunc.py

@@ -3,7 +3,6 @@
 
 """
 
-import logging
 from copy import copy
 from typing import Optional
 
@@ -16,11 +15,6 @@
 from pytensor.graph.fg import FunctionGraph
 
 
-_logger = logging.getLogger("pytensor.compile.function.pfunc")
-
-__docformat__ = "restructuredtext en"
-
-
 def rebuild_collect_shared(
     outputs,
     inputs=None,
@@ -78,10 +72,12 @@ def rebuild_collect_shared(
     shared_inputs = []
 
     def clone_v_get_shared_updates(v, copy_inputs_over):
-        """
-        Clones a variable and its inputs recursively until all are in clone_d.
-        Also appends all shared variables met along the way to shared inputs,
-        and their default_update (if applicable) to update_d and update_expr.
+        r"""Clones a variable and its inputs recursively until all are in `clone_d`.
+
+        Also, it appends all `SharedVariable`\s met along the way to
+        `shared_inputs` and their corresponding
+        `SharedVariable.default_update`\s (when applicable) to `update_d` and
+        `update_expr`.
 
         """
         # this co-recurses with clone_a
@@ -103,7 +99,7 @@ def clone_v_get_shared_updates(v, copy_inputs_over):
         elif isinstance(v, SharedVariable):
             if v not in shared_inputs:
                 shared_inputs.append(v)
-            if hasattr(v, "default_update"):
+            if v.default_update is not None:
                 # Check that v should not be excluded from the default
                 # updates list
                 if no_default_updates is False or (
@@ -419,22 +415,24 @@ def construct_pfunc_ins_and_outs(
         givens = []
 
     if not isinstance(params, (list, tuple)):
-        raise Exception("in pfunc() the first argument must be a list or a tuple")
+        raise TypeError("The `params` argument must be a list or a tuple")
 
     if not isinstance(no_default_updates, bool) and not isinstance(
         no_default_updates, list
     ):
-        raise TypeError("no_default_update should be either a boolean or a list")
+        raise TypeError("The `no_default_update` argument must be a boolean or list")
 
-    if len(updates) > 0 and any(
-        isinstance(v, Variable) for v in iter_over_pairs(updates)
+    if len(updates) > 0 and not all(
+        isinstance(pair, (tuple, list))
+        and len(pair) == 2
+        and isinstance(pair[0], Variable)
+        for pair in iter_over_pairs(updates)
     ):
-        raise ValueError(
-            "The updates parameter must be an OrderedDict/dict or a list of "
-            "lists/tuples with 2 elements"
+        raise TypeError(
+            "The `updates` parameter must be an ordered mapping or a list of pairs"
         )
 
-    # transform params into pytensor.compile.In objects.
+    # Transform params into pytensor.compile.In objects.
     inputs = [
         _pfunc_param_to_in(p, allow_downcast=allow_input_downcast) for p in params
     ]

pytensor/compile/function/types.py

@@ -32,6 +32,7 @@
 
 
 if TYPE_CHECKING:
+    from pytensor.compile.mode import Mode
     from pytensor.link.vm import VM
 
 
@@ -1391,16 +1392,24 @@ def check_unused_inputs(inputs, outputs, on_unused_input):
 
     @staticmethod
     def prepare_fgraph(
-        inputs, outputs, additional_outputs, fgraph, rewriter, linker, profile
+        inputs,
+        outputs,
+        additional_outputs,
+        fgraph: FunctionGraph,
+        mode: "Mode",
+        profile,
     ):
 
+        rewriter = mode.optimizer
+
         try:
             start_rewriter = time.perf_counter()
 
             rewriter_profile = None
             rewrite_time = None
 
             with config.change_flags(
+                mode=mode,
                 compute_test_value=config.compute_test_value_opt,
                 traceback__limit=config.traceback__compile_limit,
             ):
@@ -1440,7 +1449,7 @@ def prepare_fgraph(
                     stacklevel=3,
                 )
 
-        if not hasattr(linker, "accept"):
+        if not hasattr(mode.linker, "accept"):
             raise ValueError(
                 "'linker' parameter of FunctionMaker should be "
                 f"a Linker with an accept method or one of {list(pytensor.compile.mode.predefined_linkers.keys())}"
@@ -1511,12 +1520,8 @@ def __init__(
 
         self.fgraph = fgraph
 
-        rewriter, linker = mode.optimizer, copy.copy(mode.linker)
-
         if not no_fgraph_prep:
-            self.prepare_fgraph(
-                inputs, outputs, found_updates, fgraph, rewriter, linker, profile
-            )
+            self.prepare_fgraph(inputs, outputs, found_updates, fgraph, mode, profile)
 
         assert len(fgraph.outputs) == len(outputs + found_updates)
 
@@ -1528,6 +1533,8 @@ def __init__(
             if not spec.borrow
         ]
 
+        linker = copy.copy(mode.linker)
+
         if no_borrow:
             self.linker = linker.accept(
                 fgraph,