Make PyTensor compatible with numpy 2.0

.github/workflows/test.yml

@@ -65,7 +65,7 @@ jobs:
       - uses: pre-commit/action@v3.0.1
 
   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}}

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", "T20"]
+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

pytensor/link/c/basic.py

@@ -10,8 +10,6 @@
 from io import StringIO
 from typing import TYPE_CHECKING, Any, Optional
 
-import numpy as np
-
 from pytensor.compile.compilelock import lock_ctx
 from pytensor.configdefaults import config
 from pytensor.graph.basic import (
@@ -33,6 +31,7 @@
 from pytensor.link.c.cmodule import get_module_cache as _get_module_cache
 from pytensor.link.c.interface import CLinkerObject, CLinkerOp, CLinkerType
 from pytensor.link.utils import gc_helper, map_storage, raise_with_op, streamline
+from pytensor.npy_2_compat import ndarray_c_version
 from pytensor.utils import difference, uniq
 
 
@@ -1367,10 +1366,6 @@ def cmodule_key_(
 
         # We must always add the numpy ABI version here as
         # DynamicModule always add the include <numpy/arrayobject.h>
-        if np.lib.NumpyVersion(np.__version__) < "1.16.0a":
-            ndarray_c_version = np.core.multiarray._get_ndarray_c_version()
-        else:
-            ndarray_c_version = np.core._multiarray_umath._get_ndarray_c_version()
         sig.append(f"NPY_ABI_VERSION=0x{ndarray_c_version:X}")
         if c_compiler:
             sig.append("c_compiler_str=" + c_compiler.version_str())

pytensor/link/c/c_code/lazylinker_c.c

@@ -5,9 +5,6 @@
 
 #if PY_VERSION_HEX >= 0x03000000
 #include "numpy/npy_3kcompat.h"
-#define PyCObject_AsVoidPtr NpyCapsule_AsVoidPtr
-#define PyCObject_GetDesc NpyCapsule_GetDesc
-#define PyCObject_Check NpyCapsule_Check
 #endif
 
 #ifndef Py_TYPE
@@ -323,9 +320,9 @@ static int CLazyLinker_init(CLazyLinker *self, PyObject *args, PyObject *kwds) {
       if (PyObject_HasAttrString(thunk, "cthunk")) {
         PyObject *cthunk = PyObject_GetAttrString(thunk, "cthunk");
         // new reference
-        assert(cthunk && PyCObject_Check(cthunk));
-        self->thunk_cptr_fn[i] = PyCObject_AsVoidPtr(cthunk);
-        self->thunk_cptr_data[i] = PyCObject_GetDesc(cthunk);
+        assert(cthunk && NpyCapsule_Check(cthunk));
+        self->thunk_cptr_fn[i] = NpyCapsule_AsVoidPtr(cthunk);
+        self->thunk_cptr_data[i] = NpyCapsule_GetDesc(cthunk);
         Py_DECREF(cthunk);
         // cthunk is kept alive by membership in self->thunks
       }
@@ -487,8 +484,8 @@ static PyObject *pycall(CLazyLinker *self, Py_ssize_t node_idx, int verbose) {
       PyList_SetItem(self->call_times, node_idx,
                      PyFloat_FromDouble(t1 - t0 + ti));
       PyObject *count = PyList_GetItem(self->call_counts, node_idx);
-      long icount = PyInt_AsLong(count);
-      PyList_SetItem(self->call_counts, node_idx, PyInt_FromLong(icount + 1));
+      long icount = PyLong_AsLong(count);
+      PyList_SetItem(self->call_counts, node_idx, PyLong_FromLong(icount + 1));
     }
   } else {
     if (verbose) {
@@ -512,8 +509,8 @@ static int c_call(CLazyLinker *self, Py_ssize_t node_idx, int verbose) {
     PyList_SetItem(self->call_times, node_idx,
                    PyFloat_FromDouble(t1 - t0 + ti));
     PyObject *count = PyList_GetItem(self->call_counts, node_idx);
-    long icount = PyInt_AsLong(count);
-    PyList_SetItem(self->call_counts, node_idx, PyInt_FromLong(icount + 1));
+    long icount = PyLong_AsLong(count);
+    PyList_SetItem(self->call_counts, node_idx, PyLong_FromLong(icount + 1));
   } else {
     err = fn(self->thunk_cptr_data[node_idx]);
   }
@@ -774,20 +771,20 @@ static PyObject *CLazyLinker_call(PyObject *_self, PyObject *args,
       output_subset = (char *)calloc(self->n_output_vars, sizeof(char));
       for (int it = 0; it < output_subset_size; ++it) {
         PyObject *elem = PyList_GetItem(output_subset_ptr, it);
-        if (!PyInt_Check(elem)) {
+        if (!PyLong_Check(elem)) {
           err = 1;
           PyErr_SetString(PyExc_RuntimeError,
                           "Some elements of output_subset list are not int");
         }
-        output_subset[PyInt_AsLong(elem)] = 1;
+        output_subset[PyLong_AsLong(elem)] = 1;
       }
     }
   }
 
   self->position_of_error = -1;
   // create constants used to fill the var_compute_cells
-  PyObject *one = PyInt_FromLong(1);
-  PyObject *zero = PyInt_FromLong(0);
+  PyObject *one = PyLong_FromLong(1);
+  PyObject *zero = PyLong_FromLong(0);
 
   // pre-allocate our return value
   Py_INCREF(Py_None);
@@ -942,11 +939,8 @@ static PyMemberDef CLazyLinker_members[] = {
 };
 
 static PyTypeObject lazylinker_ext_CLazyLinkerType = {
-#if defined(NPY_PY3K)
     PyVarObject_HEAD_INIT(NULL, 0)
-#else
-    PyObject_HEAD_INIT(NULL) 0, /*ob_size*/
-#endif
+
         "lazylinker_ext.CLazyLinker",         /*tp_name*/
     sizeof(CLazyLinker),                      /*tp_basicsize*/
     0,                                        /*tp_itemsize*/
@@ -987,7 +981,7 @@ static PyTypeObject lazylinker_ext_CLazyLinkerType = {
 };
 
 static PyObject *get_version(PyObject *dummy, PyObject *args) {
-  PyObject *result = PyFloat_FromDouble(0.212);
+  PyObject *result = PyFloat_FromDouble(0.3);
   return result;
 }
 
@@ -996,7 +990,7 @@ static PyMethodDef lazylinker_ext_methods[] = {
     {NULL, NULL, 0, NULL} /* Sentinel */
 };
 
-#if defined(NPY_PY3K)
+
 static struct PyModuleDef moduledef = {PyModuleDef_HEAD_INIT,
                                        "lazylinker_ext",
                                        NULL,
@@ -1006,28 +1000,19 @@ static struct PyModuleDef moduledef = {PyModuleDef_HEAD_INIT,
                                        NULL,
                                        NULL,
                                        NULL};
-#endif
-#if defined(NPY_PY3K)
-#define RETVAL m
+
 PyMODINIT_FUNC PyInit_lazylinker_ext(void) {
-#else
-#define RETVAL
-PyMODINIT_FUNC initlazylinker_ext(void) {
-#endif
+
   PyObject *m;
 
   lazylinker_ext_CLazyLinkerType.tp_new = PyType_GenericNew;
   if (PyType_Ready(&lazylinker_ext_CLazyLinkerType) < 0)
-    return RETVAL;
-#if defined(NPY_PY3K)
+    return NULL;
+
   m = PyModule_Create(&moduledef);
-#else
-  m = Py_InitModule3("lazylinker_ext", lazylinker_ext_methods,
-                     "Example module that creates an extension type.");
-#endif
   Py_INCREF(&lazylinker_ext_CLazyLinkerType);
   PyModule_AddObject(m, "CLazyLinker",
                      (PyObject *)&lazylinker_ext_CLazyLinkerType);
 
-  return RETVAL;
+  return m;
 }

pytensor/link/c/c_code/pytensor_mod_helper.h

@@ -18,14 +18,8 @@
 #define PYTENSOR_EXTERN
 #endif
 
-#if PY_MAJOR_VERSION < 3
-#define PYTENSOR_RTYPE void
-#else
-#define PYTENSOR_RTYPE PyObject *
-#endif
-
 /* We need to redefine PyMODINIT_FUNC to add MOD_PUBLIC in the middle */
 #undef PyMODINIT_FUNC
-#define PyMODINIT_FUNC PYTENSOR_EXTERN MOD_PUBLIC PYTENSOR_RTYPE
+#define PyMODINIT_FUNC PYTENSOR_EXTERN MOD_PUBLIC PyObject *
 
 #endif

pytensor/link/c/interface.py

@@ -1,7 +1,7 @@
 import typing
 import warnings
 from abc import abstractmethod
-from collections.abc import Callable
+from collections.abc import Callable, Hashable
 from typing import Optional
 
 from pytensor.graph.basic import Apply, Constant
@@ -155,7 +155,7 @@ def c_init_code(self, **kwargs) -> list[str]:
         """Return a list of code snippets to be inserted in module initialization."""
         return []
 
-    def c_code_cache_version(self) -> tuple[int, ...]:
+    def c_code_cache_version(self) -> tuple[Hashable, ...]:
         """Return a tuple of integers indicating the version of this `Op`.
 
         An empty tuple indicates an "unversioned" `Op` that will not be cached
@@ -223,7 +223,7 @@ def c_code(
         """
         raise NotImplementedError()
 
-    def c_code_cache_version_apply(self, node: Apply) -> tuple[int, ...]:
+    def c_code_cache_version_apply(self, node: Apply) -> tuple[Hashable, ...]:
         """Return a tuple of integers indicating the version of this `Op`.
 
         An empty tuple indicates an "unversioned" `Op` that will not be

pytensor/link/c/lazylinker_c.py

@@ -14,7 +14,7 @@
 _logger = logging.getLogger(__file__)
 
 force_compile = False
-version = 0.212  # must match constant returned in function get_version()
+version = 0.3  # must match constant returned in function get_version()
 lazylinker_ext: ModuleType | None = None
 
 

pytensor/link/jax/dispatch/random.py

@@ -56,7 +56,7 @@ def assert_size_argument_jax_compatible(node):
 
 @jax_typify.register(Generator)
 def jax_typify_Generator(rng, **kwargs):
-    state = rng.__getstate__()
+    state = rng.bit_generator.state
     state["bit_generator"] = numpy_bit_gens[state["bit_generator"]]
 
     # XXX: Is this a reasonable approach?

pytensor/link/numba/dispatch/elemwise.py

@@ -4,7 +4,6 @@
 import numba
 import numpy as np
 from numba.core.extending import overload
-from numpy.core.numeric import normalize_axis_index, normalize_axis_tuple
 
 from pytensor.graph.op import Op
 from pytensor.link.numba.dispatch import basic as numba_basic
@@ -19,6 +18,7 @@
     store_core_outputs,
 )
 from pytensor.link.utils import compile_function_src
+from pytensor.npy_2_compat import normalize_axis_index, normalize_axis_tuple
 from pytensor.scalar.basic import (
     AND,
     OR,

pytensor/link/numba/dispatch/random.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from copy import copy
+from copy import copy, deepcopy
 from functools import singledispatch
 from textwrap import dedent
 
@@ -34,7 +34,7 @@
     def impl(rng):
         # TODO: Open issue on Numba?
         with numba.objmode(new_rng=types.npy_rng):
-            new_rng = copy(rng)
+            new_rng = deepcopy(rng)
 
         return new_rng