Update numpy deprecated imports

- replaced np.AxisError with np.exceptions.AxisError
- the `numpy.core` submodule has been renamed to `numpy._core`
- some parts of `numpy.core` have been moved to `numpy.lib.array_utils`

Except for `AxisError`, the updated imports are conditional on
the version of numpy, so the imports should work for numpy >= 1.26.

The conditional imports have been added to `npy_2_compat.py`, so the
imports elsewhere are unconditonal.

Author: ricardoV94

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/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/npy_2_compat.py

@@ -0,0 +1,275 @@
+from textwrap import dedent
+
+import numpy as np
+
+
+# Conditional numpy imports for numpy 1.26 and 2.x compatibility
+try:
+    from numpy.lib.array_utils import normalize_axis_index, normalize_axis_tuple
+except ModuleNotFoundError:
+    # numpy < 2.0
+    from numpy.core.multiarray import normalize_axis_index  # type: ignore[no-redef]
+    from numpy.core.numeric import normalize_axis_tuple  # type: ignore[no-redef]
+
+
+try:
+    from numpy._core.einsumfunc import (  # type: ignore[attr-defined]
+        _find_contraction,
+        _parse_einsum_input,
+    )
+except ModuleNotFoundError:
+    from numpy.core.einsumfunc import (  # type: ignore[no-redef]
+        _find_contraction,
+        _parse_einsum_input,
+    )
+
+
+# suppress linting warning by "using" the imports here:
+__all__ = [
+    "_find_contraction",
+    "_parse_einsum_input",
+    "normalize_axis_index",
+    "normalize_axis_tuple",
+]
+
+
+numpy_version_tuple = tuple(int(n) for n in np.__version__.split(".")[:2])
+numpy_version = np.lib.NumpyVersion(
+    np.__version__
+)  # used to compare with version strings, e.g. numpy_version < "1.16.0"
+using_numpy_2 = numpy_version >= "2.0.0rc1"
+
+
+if using_numpy_2:
+    ndarray_c_version = np._core._multiarray_umath._get_ndarray_c_version()
+else:
+    ndarray_c_version = np.core._multiarray_umath._get_ndarray_c_version()  # type: ignore[attr-defined]
+
+
+if using_numpy_2:
+    UintOverflowError = OverflowError
+else:
+    UintOverflowError = TypeError
+
+
+def npy_2_compat_header() -> str:
+    """Compatibility header that Numpy suggests is vendored with code that uses Numpy < 2.0 and Numpy 2.x"""
+    return dedent("""
+    #ifndef NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_
+    #define NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_
+
+
+    /*
+     * This header is meant to be included by downstream directly for 1.x compat.
+     * In that case we need to ensure that users first included the full headers
+     * and not just `ndarraytypes.h`.
+     */
+
+    #ifndef NPY_FEATURE_VERSION
+      #error "The NumPy 2 compat header requires `import_array()` for which "  \\
+             "the `ndarraytypes.h` header include is not sufficient.  Please "  \\
+             "include it after `numpy/ndarrayobject.h` or similar."  \\
+             "" \\
+             "To simplify inclusion, you may use `PyArray_ImportNumPy()` " \\
+             "which is defined in the compat header and is lightweight (can be)."
+    #endif
+
+    #if NPY_ABI_VERSION < 0x02000000
+      /*
+       * Define 2.0 feature version as it is needed below to decide whether we
+       * compile for both 1.x and 2.x (defining it gaurantees 1.x only).
+       */
+      #define NPY_2_0_API_VERSION 0x00000012
+      /*
+       * If we are compiling with NumPy 1.x, PyArray_RUNTIME_VERSION so we
+       * pretend the `PyArray_RUNTIME_VERSION` is `NPY_FEATURE_VERSION`.
+       * This allows downstream to use `PyArray_RUNTIME_VERSION` if they need to.
+       */
+      #define PyArray_RUNTIME_VERSION NPY_FEATURE_VERSION
+      /* Compiling on NumPy 1.x where these are the same: */
+      #define PyArray_DescrProto PyArray_Descr
+    #endif
+
+
+    /*
+     * Define a better way to call `_import_array()` to simplify backporting as
+     * we now require imports more often (necessary to make ABI flexible).
+     */
+    #ifdef import_array1
+
+    static inline int
+    PyArray_ImportNumPyAPI()
+    {
+        if (NPY_UNLIKELY(PyArray_API == NULL)) {
+            import_array1(-1);
+        }
+        return 0;
+    }
+
+    #endif  /* import_array1 */
+
+
+    /*
+     * NPY_DEFAULT_INT
+     *
+     * The default integer has changed, `NPY_DEFAULT_INT` is available at runtime
+     * for use as type number, e.g. `PyArray_DescrFromType(NPY_DEFAULT_INT)`.
+     *
+     * NPY_RAVEL_AXIS
+     *
+     * This was introduced in NumPy 2.0 to allow indicating that an axis should be
+     * raveled in an operation. Before NumPy 2.0, NPY_MAXDIMS was used for this purpose.
+     *
+     * NPY_MAXDIMS
+     *
+     * A constant indicating the maximum number dimensions allowed when creating
+     * an ndarray.
+     *
+     * NPY_NTYPES_LEGACY
+     *
+     * The number of built-in NumPy dtypes.
+     */
+    #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION
+        #define NPY_DEFAULT_INT NPY_INTP
+        #define NPY_RAVEL_AXIS NPY_MIN_INT
+        #define NPY_MAXARGS 64
+
+    #elif NPY_ABI_VERSION < 0x02000000
+        #define NPY_DEFAULT_INT NPY_LONG
+        #define NPY_RAVEL_AXIS 32
+        #define NPY_MAXARGS 32
+
+        /* Aliases of 2.x names to 1.x only equivalent names */
+        #define NPY_NTYPES NPY_NTYPES_LEGACY
+        #define PyArray_DescrProto PyArray_Descr
+        #define _PyArray_LegacyDescr PyArray_Descr
+        /* NumPy 2 definition always works, but add it for 1.x only */
+        #define PyDataType_ISLEGACY(dtype) (1)
+    #else
+        #define NPY_DEFAULT_INT  \\
+            (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? NPY_INTP : NPY_LONG)
+        #define NPY_RAVEL_AXIS  \\
+            (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? -1 : 32)
+        #define NPY_MAXARGS  \\
+            (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION ? 64 : 32)
+    #endif
+
+
+    /*
+     * Access inline functions for descriptor fields.  Except for the first
+     * few fields, these needed to be moved (elsize, alignment) for
+     * additional space.  Or they are descriptor specific and are not generally
+     * available anymore (metadata, c_metadata, subarray, names, fields).
+     *
+     * Most of these are defined via the `DESCR_ACCESSOR` macro helper.
+     */
+    #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION || NPY_ABI_VERSION < 0x02000000
+        /* Compiling for 1.x or 2.x only, direct field access is OK: */
+
+        static inline void
+        PyDataType_SET_ELSIZE(PyArray_Descr *dtype, npy_intp size)
+        {
+            dtype->elsize = size;
+        }
+
+        static inline npy_uint64
+        PyDataType_FLAGS(const PyArray_Descr *dtype)
+        {
+        #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION
+            return dtype->flags;
+        #else
+            return (unsigned char)dtype->flags;  /* Need unsigned cast on 1.x */
+        #endif
+        }
+
+        #define DESCR_ACCESSOR(FIELD, field, type, legacy_only)    \\
+            static inline type                                     \\
+            PyDataType_##FIELD(const PyArray_Descr *dtype) {       \\
+                if (legacy_only && !PyDataType_ISLEGACY(dtype)) {  \\
+                    return (type)0;                                \\
+                }                                                  \\
+                return ((_PyArray_LegacyDescr *)dtype)->field;     \\
+            }
+    #else  /* compiling for both 1.x and 2.x */
+
+        static inline void
+        PyDataType_SET_ELSIZE(PyArray_Descr *dtype, npy_intp size)
+        {
+            if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) {
+                ((_PyArray_DescrNumPy2 *)dtype)->elsize = size;
+            }
+            else {
+                ((PyArray_DescrProto *)dtype)->elsize = (int)size;
+            }
+        }
+
+        static inline npy_uint64
+        PyDataType_FLAGS(const PyArray_Descr *dtype)
+        {
+            if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) {
+                return ((_PyArray_DescrNumPy2 *)dtype)->flags;
+            }
+            else {
+                return (unsigned char)((PyArray_DescrProto *)dtype)->flags;
+            }
+        }
+
+        /* Cast to LegacyDescr always fine but needed when `legacy_only` */
+        #define DESCR_ACCESSOR(FIELD, field, type, legacy_only)        \\
+            static inline type                                         \\
+            PyDataType_##FIELD(const PyArray_Descr *dtype) {           \\
+                if (legacy_only && !PyDataType_ISLEGACY(dtype)) {      \\
+                    return (type)0;                                    \\
+                }                                                      \\
+                if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) {  \\
+                    return ((_PyArray_LegacyDescr *)dtype)->field;     \\
+                }                                                      \\
+                else {                                                 \\
+                    return ((PyArray_DescrProto *)dtype)->field;       \\
+                }                                                      \\
+            }
+    #endif
+
+    DESCR_ACCESSOR(ELSIZE, elsize, npy_intp, 0)
+    DESCR_ACCESSOR(ALIGNMENT, alignment, npy_intp, 0)
+    DESCR_ACCESSOR(METADATA, metadata, PyObject *, 1)
+    DESCR_ACCESSOR(SUBARRAY, subarray, PyArray_ArrayDescr *, 1)
+    DESCR_ACCESSOR(NAMES, names, PyObject *, 1)
+    DESCR_ACCESSOR(FIELDS, fields, PyObject *, 1)
+    DESCR_ACCESSOR(C_METADATA, c_metadata, NpyAuxData *, 1)
+
+    #undef DESCR_ACCESSOR
+
+
+    #if !(defined(NPY_INTERNAL_BUILD) && NPY_INTERNAL_BUILD)
+    #if NPY_FEATURE_VERSION >= NPY_2_0_API_VERSION
+        static inline PyArray_ArrFuncs *
+        PyDataType_GetArrFuncs(const PyArray_Descr *descr)
+        {
+            return _PyDataType_GetArrFuncs(descr);
+        }
+    #elif NPY_ABI_VERSION < 0x02000000
+        static inline PyArray_ArrFuncs *
+        PyDataType_GetArrFuncs(const PyArray_Descr *descr)
+        {
+            return descr->f;
+        }
+    #else
+        static inline PyArray_ArrFuncs *
+        PyDataType_GetArrFuncs(const PyArray_Descr *descr)
+        {
+            if (PyArray_RUNTIME_VERSION >= NPY_2_0_API_VERSION) {
+                return _PyDataType_GetArrFuncs(descr);
+            }
+            else {
+                return ((PyArray_DescrProto *)descr)->f;
+            }
+        }
+    #endif
+
+
+    #endif  /* not internal build */
+
+    #endif  /* NUMPY_CORE_INCLUDE_NUMPY_NPY_2_COMPAT_H_ */
+
+    """)

pytensor/tensor/__init__.py

@@ -123,7 +123,7 @@ def _get_vector_length_Constant(op: Op | Variable, var: Constant) -> int:
 
 # isort: on
 # Allow accessing numpy constants from pytensor.tensor
-from numpy import e, euler_gamma, inf, infty, nan, newaxis, pi
+from numpy import e, euler_gamma, inf, nan, newaxis, pi
 
 from pytensor.tensor.basic import *
 from pytensor.tensor.blas import batched_dot, batched_tensordot

pytensor/tensor/basic.py

@@ -14,8 +14,7 @@
 from typing import cast as type_cast
 
 import numpy as np
-from numpy.core.multiarray import normalize_axis_index
-from numpy.core.numeric import normalize_axis_tuple
+from numpy.exceptions import AxisError
 
 import pytensor
 import pytensor.scalar.sharedvar
@@ -32,6 +31,7 @@
 from pytensor.graph.type import HasShape, Type
 from pytensor.link.c.op import COp
 from pytensor.link.c.params_type import ParamsType
+from pytensor.npy_2_compat import normalize_axis_index, normalize_axis_tuple
 from pytensor.printing import Printer, min_informative_str, pprint, set_precedence
 from pytensor.raise_op import CheckAndRaise
 from pytensor.scalar import int32
@@ -228,7 +228,7 @@ def constant(x, name=None, ndim=None, dtype=None) -> TensorConstant:
         elif x_.ndim > ndim:
             try:
                 x_ = np.squeeze(x_, axis=tuple(range(x_.ndim - ndim)))
-            except np.AxisError:
+            except AxisError:
                 raise ValueError(
                     f"ndarray could not be cast to constant with {int(ndim)} dimensions"
                 )
@@ -4405,7 +4405,7 @@ def expand_dims(a: np.ndarray | TensorVariable, axis: Sequence[int]) -> TensorVa
         axis = (axis,)
 
     out_ndim = len(axis) + a.ndim
-    axis = np.core.numeric.normalize_axis_tuple(axis, out_ndim)
+    axis = normalize_axis_tuple(axis, out_ndim)
 
     if not axis:
         return a

pytensor/tensor/conv/abstract_conv.py

@@ -8,6 +8,7 @@
 from math import gcd
 
 import numpy as np
+from numpy.exceptions import ComplexWarning
 
 
 try:
@@ -2338,7 +2339,7 @@ def conv(
             bval = _bvalfromboundary("fill")
 
             with warnings.catch_warnings():
-                warnings.simplefilter("ignore", np.ComplexWarning)
+                warnings.simplefilter("ignore", ComplexWarning)
                 for b in range(img.shape[0]):
                     for g in range(self.num_groups):
                         for n in range(output_channel_offset):