Numpy 2 updates

.github/workflows/pypi.yml

@@ -189,5 +189,5 @@ jobs:
           name: universal_wheel
           path: dist
 
-      - uses: pypa/[email protected].2
+      - uses: pypa/[email protected].4
         # Implicitly attests that the packages were uploaded in the context of this workflow.

.github/workflows/test.yml

@@ -65,7 +65,7 @@ jobs:
       - uses: pre-commit/[email protected]
 
   test:
-    name: "${{ matrix.os }} test py${{ matrix.python-version }} : fast-compile ${{ matrix.fast-compile }} : float32 ${{ matrix.float32 }} : ${{ matrix.part }}"
+    name: "${{ matrix.os }} test py${{ matrix.python-version }} numpy${{ matrix.numpy-version }} : fast-compile ${{ matrix.fast-compile }} : float32 ${{ matrix.float32 }} : ${{ matrix.part }}"
     needs:
       - changes
       - style
@@ -76,6 +76,7 @@ jobs:
       matrix:
         os: ["ubuntu-latest"]
         python-version: ["3.10", "3.12"]
+        numpy-version: ["~=1.26.0", ">=2.0"]
         fast-compile: [0, 1]
         float32: [0, 1]
         install-numba: [0]
@@ -105,45 +106,68 @@ jobs:
             float32: 1
           - part: "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link"
             fast-compile: 1
+          - numpy-version: "~=1.26.0"
+            fast-compile: 1
+          - numpy-version: "~=1.26.0"
+            float32: 1
+          - numpy-version: "~=1.26.0"
+            python-version: "3.12"
+          - numpy-version: "~=1.26.0"
+            part: "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link"
         include:
           - install-numba: 1
             os: "ubuntu-latest"
             python-version: "3.10"
+            numpy-version: "~=2.1.0"
             fast-compile: 0
             float32: 0
             part: "tests/link/numba"
           - install-numba: 1
             os: "ubuntu-latest"
             python-version: "3.12"
+            numpy-version: "~=2.1.0"
             fast-compile: 0
             float32: 0
             part: "tests/link/numba"
           - install-jax: 1
             os: "ubuntu-latest"
             python-version: "3.10"
+            numpy-version: ">=2.0"
             fast-compile: 0
             float32: 0
             part: "tests/link/jax"
           - install-jax: 1
             os: "ubuntu-latest"
             python-version: "3.12"
+            numpy-version: ">=2.0"
             fast-compile: 0
             float32: 0
             part: "tests/link/jax"
           - install-torch: 1
             os: "ubuntu-latest"
             python-version: "3.10"
+            numpy-version: ">=2.0"
             fast-compile: 0
             float32: 0
             part: "tests/link/pytorch"
           - os: macos-15
             python-version: "3.12"
+            numpy-version: ">=2.0"
             fast-compile: 0
             float32: 0
             install-numba: 0
             install-jax: 0
             install-torch: 0
             part: "tests/tensor/test_blas.py tests/tensor/test_elemwise.py tests/tensor/test_math_scipy.py"
+          - os: "ubuntu-latest"
+            python-version: "3.10"
+            numpy-version: "~=1.26.0"
+            fast-compile: 0
+            float32: 0
+            install-numba: 0
+            install-jax: 0
+            install-torch: 0
+            part: "tests/tensor/test_math.py"
 
     steps:
       - uses: actions/checkout@v4
@@ -174,9 +198,9 @@ jobs:
         run: |
 
           if [[ $OS == "macos-15" ]]; then
-            micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" numpy scipy pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock libblas=*=*accelerate;
+            micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" "numpy${NUMPY_VERSION}" scipy pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock libblas=*=*accelerate;
           else
-            micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock;
+            micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl "numpy${NUMPY_VERSION}" scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock;
           fi
           if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi
           if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi
@@ -193,6 +217,7 @@ jobs:
           fi
         env:
           PYTHON_VERSION: ${{ matrix.python-version }}
+          NUMPY_VERSION:  ${{ matrix.numpy-version }}
           INSTALL_NUMBA: ${{ matrix.install-numba }}
           INSTALL_JAX: ${{ matrix.install-jax }}
           INSTALL_TORCH: ${{ matrix.install-torch}}

doc/extending/op.rst

@@ -506,4 +506,3 @@ These are the function required to work with :func:`pytensor.gradient.grad`.
    the outputs) back to their corresponding shapes and return them as the
    output of the :meth:`Op.R_op` method.
 
-   :ref:`List of op with r op support <R_op_list>`.

doc/library/tensor/basic.rst

@@ -1791,5 +1791,3 @@ Gradient / Differentiation
     :members: grad
     :noindex:
 
-See the :ref:`gradient <libdoc_gradient>` page for complete documentation
-of the gradient module.

doc/tutorial/gradients.rst

@@ -86,9 +86,7 @@ of symbolic differentiation).
    ``i`` of the output list is the gradient of the first argument of
    `pt.grad` with respect to the ``i``-th element of the list given as second argument.
    The first argument of `pt.grad` has to be a scalar (a tensor
-   of size 1). For more information on the semantics of the arguments of
-   `pt.grad` and details about the implementation, see
-   :ref:`this<libdoc_gradient>` section of the library.
+   of size 1).
 
    Additional information on the inner workings of differentiation may also be
    found in the more advanced tutorial :ref:`Extending PyTensor<extending>`.
@@ -204,7 +202,21 @@ you need to do something similar to this:
 >>> f([[1, 1], [1, 1]], [[2, 2], [2, 2]], [0,1])
 array([ 2.,  2.])
 
-:ref:`List <R_op_list>` of Op that implement Rop.
+By default, the R-operator is implemented as a double application of the L_operator
+(see `reference <https://j-towns.github.io/2017/06/12/A-new-trick.html>`_).
+In most cases this should be as performant as a specialized implementation of the R-operator.
+However, PyTensor may sometimes fail to prune dead branches or fuse common expressions within composite operators,
+such as Scan and OpFromGraph, that would be more easily avoidable in a direct implentation of the R-operator.
+
+When this is a concern, it is possible to force `Rop` to use the specialized `Op.R_op` methods by passing
+`use_op_rop_implementation=True`. Note that this will fail if the graph contains `Op`s that don't implement this method.
+
+
+>>> JV = pytensor.gradient.Rop(y, W, V, use_op_rop_implementation=True)
+>>> f = pytensor.function([W, V, x], JV)
+>>> f([[1, 1], [1, 1]], [[2, 2], [2, 2]], [0,1])
+array([ 2.,  2.])
+
 
 L-operator
 ----------
@@ -234,7 +246,6 @@ array([[ 0.,  0.],
     as the input parameter, while the result of the R-operator has a shape similar
     to that of the output.
 
-   :ref:`List of op with r op support <R_op_list>`.
 
 Hessian times a Vector
 ======================

environment-osx-arm64.yml

@@ -9,7 +9,7 @@ channels:
 dependencies:
   - python=>3.10
   - compilers
-  - numpy>=1.17.0,<2
+  - numpy>=1.17.0
   - scipy>=1,<2
   - filelock>=3.15
   - etuples

environment.yml

@@ -9,7 +9,7 @@ channels:
 dependencies:
   - python>=3.10
   - compilers
-  - numpy>=1.17.0,<2
+  - numpy>=1.17.0
   - scipy>=1,<2
   - filelock>=3.15
   - etuples

pyproject.toml

@@ -48,7 +48,7 @@ keywords = [
 dependencies = [
     "setuptools>=59.0.0",
     "scipy>=1,<2",
-    "numpy>=1.17.0,<2",
+    "numpy>=1.17.0",
     "filelock>=3.15",
     "etuples",
     "logical-unification",
@@ -129,7 +129,7 @@ exclude = ["doc/", "pytensor/_version.py"]
 docstring-code-format = true
 
 [tool.ruff.lint]
-select = ["B905", "C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH", "ISC"]
+select = ["B905", "C", "E", "F", "I", "UP", "W", "RUF", "PERF", "PTH", "ISC", "T20", "NPY201"]
 ignore = ["C408", "C901", "E501", "E741", "RUF012", "PERF203", "ISC001"]
 unfixable = [
     # zip-strict: the auto-fix adds `strict=False` but we might want `strict=True` instead
@@ -144,7 +144,12 @@ lines-after-imports = 2
 # TODO: Get rid of these:
 "**/__init__.py" = ["F401", "E402", "F403"]
 "pytensor/tensor/linalg.py" = ["F403"]
-"pytensor/link/c/cmodule.py" = ["PTH"]
+"pytensor/link/c/cmodule.py" = ["PTH", "T201"]
+"pytensor/misc/elemwise_time_test.py" = ["T201"]
+"pytensor/misc/elemwise_openmp_speedup.py" = ["T201"]
+"pytensor/misc/check_duplicate_key.py" = ["T201"]
+"pytensor/misc/check_blas.py" = ["T201"]
+"pytensor/bin/pytensor_cache.py" = ["T201"]
 # For the tests we skip because `pytest.importorskip` is used:
 "tests/link/jax/test_scalar.py" = ["E402"]
 "tests/link/jax/test_tensor_basic.py" = ["E402"]
@@ -158,6 +163,8 @@ lines-after-imports = 2
 "tests/sparse/test_sp2.py" = ["E402"]
 "tests/sparse/test_utils.py" = ["E402"]
 "tests/sparse/sandbox/test_sp.py" = ["E402", "F401"]
+"tests/compile/test_monitormode.py" = ["T201"]
+"scripts/run_mypy.py" = ["T201"]
 
 
 [tool.mypy]

pytensor/breakpoint.py

@@ -108,14 +108,14 @@ def perform(self, node, inputs, output_storage):
                     f"'{self.name}' could not be casted to NumPy arrays"
                 )
 
-            print("\n")
-            print("-------------------------------------------------")
-            print(f"Conditional breakpoint '{self.name}' activated\n")
-            print("The monitored variables are stored, in order,")
-            print("in the list variable 'monitored' as NumPy arrays.\n")
-            print("Their contents can be altered and, when execution")
-            print("resumes, the updated values will be used.")
-            print("-------------------------------------------------")
+            print("\n")  # noqa: T201
+            print("-------------------------------------------------")  # noqa: T201
+            print(f"Conditional breakpoint '{self.name}' activated\n")  # noqa: T201
+            print("The monitored variables are stored, in order,")  # noqa: T201
+            print("in the list variable 'monitored' as NumPy arrays.\n")  # noqa: T201
+            print("Their contents can be altered and, when execution")  # noqa: T201
+            print("resumes, the updated values will be used.")  # noqa: T201
+            print("-------------------------------------------------")  # noqa: T201
 
             try:
                 import pudb

pytensor/compile/builders.py

@@ -340,6 +340,12 @@ def __init__(
             ``None``, this will be used as the connection_pattern for this
             :class:`Op`.
 
+        .. warning::
+
+            rop overrides is ignored when `pytensor.gradient.Rop` is called with
+            `use_op_rop_implementation=False` (default). In this case the Lop
+            is used twice to obtain a mathematically equivalent Rop.
+
         strict: bool, default False
             If true, it raises when any variables needed to compute the inner graph
             are not provided as explici inputs. This can only happen for graphs with
@@ -641,7 +647,12 @@ def _build_and_cache_rop_op(self):
             return rop_overrides
 
         eval_points = [inp_t() for inp_t in self.input_types]
-        fn_rop = partial(Rop, wrt=inner_inputs, eval_points=eval_points)
+        fn_rop = partial(
+            Rop,
+            wrt=inner_inputs,
+            eval_points=eval_points,
+            use_op_rop_implementation=True,
+        )
 
         callable_args = (inner_inputs, eval_points)
         if rop_overrides is None: