Merge branch 'main' into numba_slinalg

Author: jessegrabowski

.github/workflows/pypi.yml

@@ -22,7 +22,7 @@ jobs:
     name: Make SDist
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
           submodules: true
@@ -44,7 +44,7 @@ jobs:
           - windows-2022
           - ubuntu-20.04
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 

.github/workflows/test.yml

@@ -22,7 +22,7 @@ jobs:
     outputs:
       changes: ${{ steps.changes.outputs.src }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - uses: dorny/paths-filter@v2
@@ -54,7 +54,7 @@ jobs:
       matrix:
         python-version: ["3.9", "3.10", "3.11"]
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@v4
     - uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
@@ -116,7 +116,7 @@ jobs:
             float32: 0
             part: "tests/link/jax"
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Set up Python ${{ matrix.python-version }}
@@ -139,7 +139,7 @@ jobs:
       - name: Install dependencies
         shell: bash -l {0}
         run: |
-          mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
+          mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl "numpy<1.26" scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
           # numba-scipy downgrades the installed scipy to 1.7.3 in Python 3.9, but
           # not numpy, even though scipy 1.7 requires numpy<1.23. When installing
           # PyTensor next, pip installs a lower version of numpy via the PyPI.
@@ -187,7 +187,7 @@ jobs:
       strategy:
         fail-fast: false
       steps:
-        - uses: actions/checkout@v3
+        - uses: actions/checkout@v4
           with:
             fetch-depth: 0
         - name: Set up Python 3.9
@@ -244,7 +244,7 @@ jobs:
     needs: [changes, all-checks]
     if: ${{ needs.changes.outputs.changes == 'true' && needs.all-checks.result == 'success' }}
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
 
       - name: Set up Python
         uses: actions/setup-python@v4

environment.yml

@@ -9,7 +9,7 @@ channels:
 dependencies:
   - python
   - compilers
-  - numpy>=1.17.0
+  - numpy>=1.17.0,<1.26.0
   - scipy>=0.14
   - filelock
   - etuples

pyproject.toml

@@ -2,7 +2,7 @@
 requires = [
   "setuptools>=48.0.0",
   "cython",
-  "numpy>=1.17.0",
+  "numpy>=1.17.0,<1.26",
   "versioneer[toml]==0.28",
 ]
 build-backend = "setuptools.build_meta"
@@ -52,7 +52,7 @@ keywords = [
 dependencies = [
     "setuptools>=48.0.0",
     "scipy>=0.14",
-    "numpy>=1.17.0",
+    "numpy>=1.17.0,<1.26",
     "filelock",
     "etuples",
     "logical-unification",

pytensor/compile/mode.py

@@ -251,13 +251,14 @@ def apply(self, fgraph):
 # especially constant merge
 optdb.register("merge2", MergeOptimizer(), "fast_run", "merge", position=49)
 
+optdb.register("py_only", EquilibriumDB(), "fast_compile", position=49.1)
+
 optdb.register(
     "add_destroy_handler", AddDestroyHandler(), "fast_run", "inplace", position=49.5
 )
 
 # final pass just to make sure
 optdb.register("merge3", MergeOptimizer(), "fast_run", "merge", position=100)
-optdb.register("py_only", EquilibriumDB(), "fast_compile", position=100)
 
 _tags: Union[Tuple[str, str], Tuple]
 
@@ -463,7 +464,10 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
 )
 NUMBA = Mode(
     NumbaLinker(),
-    RewriteDatabaseQuery(include=["fast_run"], exclude=["cxx_only", "BlasOpt"]),
+    RewriteDatabaseQuery(
+        include=["fast_run"],
+        exclude=["cxx_only", "BlasOpt", "local_careduce_fusion"],
+    ),
 )
 
 

pytensor/gradient.py

@@ -196,12 +196,13 @@ def Rop(
 
     Returns
     -------
+    :class:`~pytensor.graph.basic.Variable` or list/tuple of Variables
         A symbolic expression such obeying
         ``R_op[i] = sum_j (d f[i] / d wrt[j]) eval_point[j]``,
         where the indices in that expression are magic multidimensional
         indices that specify both the position within a list and all
         coordinates of the tensor elements.
-        If `wrt` is a list/tuple, then return a list/tuple with the results.
+        If `f` is a list/tuple, then return a list/tuple with the results.
     """
 
     if not isinstance(wrt, (list, tuple)):
@@ -384,6 +385,7 @@ def Lop(
 
     Returns
     -------
+    :class:`~pytensor.graph.basic.Variable` or list/tuple of Variables
         A symbolic expression satisfying
         ``L_op[i] = sum_i (d f[i] / d wrt[j]) eval_point[i]``
         where the indices in that expression are magic multidimensional
@@ -481,10 +483,10 @@ def grad(
 
     Returns
     -------
+    :class:`~pytensor.graph.basic.Variable` or list/tuple of Variables
         A symbolic expression for the gradient of `cost` with respect to each
         of the `wrt` terms.  If an element of `wrt` is not differentiable with
         respect to the output, then a zero variable is returned.
-
     """
     t0 = time.perf_counter()
 
@@ -701,7 +703,6 @@ def subgraph_grad(wrt, end, start=None, cost=None, details=False):
 
     Parameters
     ----------
-
     wrt : list of variables
         Gradients are computed with respect to `wrt`.
 
@@ -876,7 +877,6 @@ def _populate_var_to_app_to_idx(outputs, wrt, consider_constant):
 
     (A variable in consider_constant is not a function of
     anything)
-
     """
 
     # Validate and format consider_constant
@@ -1035,7 +1035,6 @@ def _populate_grad_dict(var_to_app_to_idx, grad_dict, wrt, cost_name=None):
     -------
     list of Variables
         A list of gradients corresponding to `wrt`
-
     """
     # build a dict mapping node to the terms node contributes to each of
     # its inputs' gradients
@@ -1423,8 +1422,9 @@ def access_grad_cache(var):
 
 
 def _float_zeros_like(x):
-    """Like zeros_like, but forces the object to have a
-    a floating point dtype"""
+    """Like zeros_like, but forces the object to have
+    a floating point dtype
+    """
 
     rval = x.zeros_like()
 
@@ -1436,7 +1436,8 @@ def _float_zeros_like(x):
 
 def _float_ones_like(x):
     """Like ones_like, but forces the object to have a
-    floating point dtype"""
+    floating point dtype
+    """
 
     dtype = x.type.dtype
     if dtype not in pytensor.tensor.type.float_dtypes:
@@ -1613,7 +1614,6 @@ def abs_rel_errors(self, g_pt):
 
         Corresponding ndarrays in `g_pt` and `self.gf` must have the same
         shape or ValueError is raised.
-
         """
         if len(g_pt) != len(self.gf):
             raise ValueError("argument has wrong number of elements", len(g_pt))
@@ -1740,7 +1740,6 @@ def verify_grad(
     This function does not support multiple outputs. In `tests.scan.test_basic`
     there is an experimental `verify_grad` that covers that case as well by
     using random projections.
-
     """
     from pytensor.compile.function import function
     from pytensor.compile.sharedvalue import shared
@@ -2267,7 +2266,6 @@ def grad_clip(x, lower_bound, upper_bound):
     -----
     We register an opt in tensor/opt.py that remove the GradClip.
     So it have 0 cost in the forward and only do work in the grad.
-
     """
     return GradClip(lower_bound, upper_bound)(x)
 

pytensor/graph/__init__.py

@@ -9,7 +9,7 @@
     clone,
     ancestors,
 )
-from pytensor.graph.replace import clone_replace, graph_replace
+from pytensor.graph.replace import clone_replace, graph_replace, vectorize
 from pytensor.graph.op import Op
 from pytensor.graph.type import Type
 from pytensor.graph.fg import FunctionGraph

pytensor/graph/replace.py

@@ -1,8 +1,9 @@
-from functools import partial
-from typing import Iterable, Optional, Sequence, Union, cast, overload
+from functools import partial, singledispatch
+from typing import Iterable, Mapping, Optional, Sequence, Union, cast, overload
 
 from pytensor.graph.basic import Apply, Constant, Variable, truncated_graph_inputs
 from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.op import Op
 
 
 ReplaceTypes = Union[Iterable[tuple[Variable, Variable]], dict[Variable, Variable]]
@@ -198,3 +199,112 @@ def toposort_key(
         return list(fg.outputs)
     else:
         return fg.outputs[0]
+
+
+@singledispatch
+def _vectorize_node(op: Op, node: Apply, *bached_inputs) -> Apply:
+    # Default implementation is provided in pytensor.tensor.blockwise
+    raise NotImplementedError
+
+
+def vectorize_node(node: Apply, *batched_inputs) -> Apply:
+    """Returns vectorized version of node with new batched inputs."""
+    op = node.op
+    return _vectorize_node(op, node, *batched_inputs)
+
+
+@overload
+def vectorize(
+    outputs: Variable,
+    replace: Mapping[Variable, Variable],
+) -> Variable:
+    ...
+
+
+@overload
+def vectorize(
+    outputs: Sequence[Variable],
+    replace: Mapping[Variable, Variable],
+) -> Sequence[Variable]:
+    ...
+
+
+def vectorize(
+    outputs: Union[Variable, Sequence[Variable]],
+    replace: Mapping[Variable, Variable],
+) -> Union[Variable, Sequence[Variable]]:
+    """Vectorize outputs graph given mapping from old variables to expanded counterparts version.
+
+    Expanded dimensions must be on the left. Behavior is similar to the functional `numpy.vectorize`.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        import pytensor
+        import pytensor.tensor as pt
+
+        from pytensor.graph import vectorize
+
+        # Original graph
+        x = pt.vector("x")
+        y = pt.exp(x) / pt.sum(pt.exp(x))
+
+        # Vectorized graph
+        new_x = pt.matrix("new_x")
+        new_y = vectorize(y, replace={x: new_x})
+
+        fn = pytensor.function([new_x], new_y)
+        fn([[0, 1, 2], [2, 1, 0]])
+        # array([[0.09003057, 0.24472847, 0.66524096],
+        #        [0.66524096, 0.24472847, 0.09003057]])
+
+
+    .. code-block:: python
+
+        import pytensor
+        import pytensor.tensor as pt
+
+        from pytensor.graph import vectorize
+
+        # Original graph
+        x = pt.vector("x")
+        y1 = x[0]
+        y2 = x[-1]
+
+        # Vectorized graph
+        new_x = pt.matrix("new_x")
+        [new_y1, new_y2] = vectorize([y1, y2], replace={x: new_x})
+
+        fn = pytensor.function([new_x], [new_y1, new_y2])
+        fn([[-10, 0, 10], [-11, 0, 11]])
+        # [array([-10., -11.]), array([10., 11.])]
+
+    """
+    if isinstance(outputs, Sequence):
+        seq_outputs = outputs
+    else:
+        seq_outputs = [outputs]
+
+    inputs = truncated_graph_inputs(seq_outputs, ancestors_to_include=replace.keys())
+    new_inputs = [replace.get(inp, inp) for inp in inputs]
+
+    def transform(var: Variable) -> Variable:
+        if var in inputs:
+            return new_inputs[inputs.index(var)]
+
+        node = var.owner
+        batched_inputs = [transform(inp) for inp in node.inputs]
+        batched_node = vectorize_node(node, *batched_inputs)
+        batched_var = batched_node.outputs[var.owner.outputs.index(var)]
+
+        return cast(Variable, batched_var)
+
+    # TODO: MergeOptimization or node caching?
+    seq_vect_outputs = [transform(out) for out in seq_outputs]
+
+    if isinstance(outputs, Sequence):
+        return seq_vect_outputs
+    else:
+        [vect_output] = seq_vect_outputs
+        return vect_output

pytensor/ifelse.py

@@ -482,6 +482,7 @@ def cond_make_inplace(fgraph, node):
     at.basic.Alloc,
     at.elemwise.Elemwise,
     at.elemwise.DimShuffle,
+    at.blockwise.Blockwise,
 )
 
 

pytensor/link/jax/dispatch/scan.py

@@ -1,6 +1,7 @@
 import jax
 import jax.numpy as jnp
 
+from pytensor.compile.mode import JAX
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.scan.op import Scan
 
@@ -17,8 +18,8 @@ def jax_funcify_Scan(op: Scan, **kwargs):
             "Scan with MIT-MOT (gradients of scan) cannot yet be converted to JAX"
         )
 
-    # Optimize inner graph
-    rewriter = op.mode_instance.optimizer
+    # Optimize inner graph (exclude any defalut rewrites that are incompatible with JAX mode)
+    rewriter = op.mode_instance.excluding(*JAX._optimizer.exclude).optimizer
     rewriter(op.fgraph)
     scan_inner_func = jax_funcify(op.fgraph, **kwargs)