use NumpyVersion

Author: vtavana

mkl_fft/_numpy_fft.py

@@ -70,7 +70,7 @@ def _cook_nd_args(a, s=None, axes=None, invreal=False):
         shapeless = False
     s = list(s)
     if axes is None:
-        if not shapeless and np.__version__ >= "2.0":
+        if not shapeless and np.lib.NumpyVersion(np.__version__) >= "2.0.0":
             msg = (
                 "`axes` should not be `None` if `s` is not `None` "
                 "(Deprecated in NumPy 2.0). In a future version of NumPy, "
@@ -85,7 +85,7 @@ def _cook_nd_args(a, s=None, axes=None, invreal=False):
         raise ValueError("Shape and axes have different lengths.")
     if invreal and shapeless:
         s[-1] = (a.shape[axes[-1]] - 1) * 2
-    if None in s and np.__version__ >= "2.0":
+    if None in s and np.lib.NumpyVersion(np.__version__) >= "2.0.0":
         msg = (
             "Passing an array containing `None` values to `s` is "
             "deprecated in NumPy 2.0 and will raise an error in "

mkl_fft/tests/helper.py

@@ -29,5 +29,6 @@
 import pytest
 
 requires_numpy_2 = pytest.mark.skipif(
-    np.__version__ < "2.0", reason="Requires NumPy >= 2.0"
+    np.lib.NumpyVersion(np.__version__) < "2.0.0",
+    reason="Requires NumPy >= 2.0.0",
 )