Remove npy<2 compatibility for NPY_RAVEL_AXIS value

Author: Armavica

pytensor/npy_2_compat.py

@@ -16,14 +16,6 @@
     ndarray_c_version = np.core._multiarray_umath._get_ndarray_c_version()  # type: ignore[attr-defined]
 
 
-# to patch up some of the C code, we need to use these special values...
-if using_numpy_2:
-    numpy_axis_is_none_flag = np.iinfo(np.int32).min  # the value of "NPY_RAVEL_AXIS"
-else:
-    # 32 is the value used to mark axis = None in Numpy C-API prior to version 2.0
-    numpy_axis_is_none_flag = 32
-
-
 # function that replicates np.unique from numpy < 2.0
 def old_np_unique(
     arr, return_index=False, return_inverse=False, return_counts=False, axis=None

pytensor/tensor/extra_ops.py

@@ -18,11 +18,7 @@
 from pytensor.link.c.op import COp
 from pytensor.link.c.params_type import ParamsType
 from pytensor.link.c.type import EnumList, Generic
-from pytensor.npy_2_compat import (
-    npy_2_compat_header,
-    numpy_axis_is_none_flag,
-    old_np_unique,
-)
+from pytensor.npy_2_compat import npy_2_compat_header, old_np_unique
 from pytensor.raise_op import Assert
 from pytensor.scalar import int64 as int_t
 from pytensor.scalar import upcast
@@ -51,7 +47,7 @@
 from pytensor.tensor.type import TensorType, dvector, int_dtypes, integer_dtypes, vector
 from pytensor.tensor.utils import normalize_reduce_axis
 from pytensor.tensor.variable import TensorVariable
-from pytensor.utils import LOCAL_BITWIDTH, PYTHON_INT_BITWIDTH
+from pytensor.utils import LOCAL_BITWIDTH, NPY_RAVEL_AXIS, PYTHON_INT_BITWIDTH
 
 
 class CpuContiguous(COp):
@@ -308,7 +304,7 @@ def __init__(self, axis: int | None = None, mode="add"):
     @property
     def c_axis(self) -> int:
         if self.axis is None:
-            return numpy_axis_is_none_flag
+            return NPY_RAVEL_AXIS
         return self.axis
 
     def make_node(self, x):

pytensor/tensor/math.py

@@ -14,7 +14,7 @@
 from pytensor.graph.replace import _vectorize_node
 from pytensor.link.c.op import COp
 from pytensor.link.c.params_type import ParamsType
-from pytensor.npy_2_compat import npy_2_compat_header, numpy_axis_is_none_flag
+from pytensor.npy_2_compat import npy_2_compat_header
 from pytensor.printing import pprint
 from pytensor.raise_op import Assert
 from pytensor.scalar.basic import BinaryScalarOp
@@ -162,7 +162,7 @@ def get_params(self, node):
             c_axis = np.int64(self.axis[0])
         else:
             # The value here doesn't matter, it won't be used
-            c_axis = numpy_axis_is_none_flag
+            c_axis = 0
         return self.params_type.get_params(c_axis=c_axis)
 
     def make_node(self, x):

pytensor/utils.py

@@ -10,6 +10,8 @@
 from functools import partial
 from pathlib import Path
 
+import numpy as np
+
 
 __all__ = [
     "get_unbound_function",
@@ -19,6 +21,7 @@
     "output_subprocess_Popen",
     "LOCAL_BITWIDTH",
     "PYTHON_INT_BITWIDTH",
+    "NPY_RAVEL_AXIS",
     "NoDuplicateOptWarningFilter",
 ]
 
@@ -46,6 +49,11 @@
 'l' denotes a C long int, and the size is expressed in bytes.
 """
 
+NPY_RAVEL_AXIS = np.iinfo(np.int32).min
+"""
+The value of the numpy C API NPY_RAVEL_AXIS.
+"""
+
 
 def __call_excepthooks(type, value, trace):
     """