update available types for ldexp

Author: Vahid Tavanashad

.github/workflows/conda-package.yml

@@ -325,10 +325,13 @@ jobs:
 
       - name: Smoke test
         shell: cmd /C CALL {0}
-        run: >-
-          conda activate mkl_umath_test && python -c "import mkl_umath, numpy as np; mkl_umath.use_in_numpy(); np.sin(np.linspace(0, 1, num=10**6));"
+        run: |
+          @ECHO ON
+          conda activate mkl_umath_test
+          python -c "import mkl_umath, numpy as np; mkl_umath.use_in_numpy(); np.sin(np.linspace(0, 1, num=10**6));"
 
       - name: Run tests
         shell: cmd /C CALL {0}
         run: |
-          conda activate mkl_umath_test && python -m pytest -v -s --pyargs ${{ env.PACKAGE_NAME }}
+          conda activate mkl_umath_test
+          python -m pytest -v -s --pyargs ${{ env.PACKAGE_NAME }}

mkl_umath/generate_umath.py

@@ -40,6 +40,7 @@
 import re
 import textwrap
 import argparse
+import numpy as np
 
 # identity objects
 Zero = "PyLong_FromLong(0)"
@@ -379,6 +380,22 @@ def english_upper(s):
 # all the function names and their corresponding ufunc signatures.  TD is
 # an object which expands a list of character codes into an array of
 # TypeDescriptions.
+
+if np.lib.NumpyVersion(np.__version__) < "2.0.0":
+    ldexp_signature = [
+        TypeDescription('f', None, 'fi', 'f'),
+        TypeDescription('f', FuncNameSuffix('long'), 'fl', 'f'),
+        TypeDescription('d', None, 'di', 'd'),
+        TypeDescription('d', FuncNameSuffix('long'), 'dl', 'd'),
+    ]
+else:
+    ldexp_signature = [
+        TypeDescription('f', None, 'fi', 'f'),
+        TypeDescription('f', FuncNameSuffix('int64'), 'f'+int64, 'f'),
+        TypeDescription('d', None, 'di', 'd'),
+        TypeDescription('d', FuncNameSuffix('int64'), 'd'+int64, 'd'),
+    ]
+
 defdict = {
 'add':
     Ufunc(2, 1, Zero,
@@ -766,12 +783,7 @@ def english_upper(s):
     Ufunc(2, 1, None,
           docstrings.get('numpy._core.umath.ldexp'),
           None,
-          [
-              TypeDescription('f', None, 'fi', 'f'),
-              TypeDescription('f', FuncNameSuffix('long'), 'fl', 'f'),
-              TypeDescription('d', None, 'di', 'd'),
-              TypeDescription('d', FuncNameSuffix('long'), 'dl', 'd'),
-          ],
+          ldexp_signature,
           ),
 'frexp' :
     Ufunc(1, 2, None,

mkl_umath/src/_patch.pyx

@@ -147,7 +147,7 @@ def use_in_numpy():
 
     Examples
     --------
-    >>> import mkl_umath, numpy as np
+    >>> import mkl_umath
     >>> mkl_umath.is_patched()
     # False
 
@@ -171,7 +171,7 @@ def restore():
 
     Examples
     --------
-    >>> import mkl_umath, numpy as np
+    >>> import mkl_umath
     >>> mkl_umath.is_patched()
     # False
 
@@ -195,7 +195,7 @@ def is_patched():
 
     Examples
     --------
-    >>> import mkl_umath, numpy as np
+    >>> import mkl_umath
     >>> mkl_umath.is_patched()
     # False
 
@@ -221,7 +221,7 @@ class mkl_umath(ContextDecorator):
 
     Examples
     --------
-    >>> import mkl_umath, numpy as np
+    >>> import mkl_umath
     >>> mkl_umath.is_patched()
     # False
 

mkl_umath/src/mkl_umath_loops.c.src

@@ -2190,6 +2190,37 @@ mkl_umath_@TYPE@_ldexp(char **args, const npy_intp *dimensions, const npy_intp *
     }
 }
 
+#ifdef USE_NUMPY_2
+void
+mkl_umath_@TYPE@_ldexp_int64(char **args, const npy_intp *dimensions, const npy_intp *steps, void *NPY_UNUSED(func))
+{
+    /*
+     * Additional loop to handle npy_long integer inputs (cf. #866, #1633).
+     * npy_long != npy_int on many 64-bit platforms, so we need this second loop
+     * to handle the default (and larger) integer types.
+     */
+    BINARY_LOOP {
+        const @type@ in1 = *(@type@ *)ip1;
+        const npy_int64 in2 = *(npy_int64 *)ip2;
+        if (((int)in2) == in2) {
+            /* Range OK */
+            *((@type@ *)op1) = ldexp@c@(in1, ((int)in2));
+        }
+        else {
+            /*
+             * Outside npy_int range -- also ldexp will overflow in this case,
+             * given that exponent has less bits than npy_int.
+             */
+            if (in2 > 0) {
+                *((@type@ *)op1) = ldexp@c@(in1, NPY_MAX_INT);
+            }
+            else {
+                *((@type@ *)op1) = ldexp@c@(in1, NPY_MIN_INT);
+            }
+        }
+    }
+}
+#else
 void
 mkl_umath_@TYPE@_ldexp_long(char **args, const npy_intp *dimensions, const npy_intp *steps, void *NPY_UNUSED(func))
 {
@@ -2219,6 +2250,7 @@ mkl_umath_@TYPE@_ldexp_long(char **args, const npy_intp *dimensions, const npy_i
         }
     }
 }
+#endif
 
 #define mkl_umath_@TYPE@_true_divide mkl_umath_@TYPE@_divide
 

mkl_umath/src/mkl_umath_loops.h.src

@@ -42,6 +42,12 @@
 #define MKL_UMATH_API
 #endif
 
+// NPY_2_0_API_VERSION 0x00000012 is defined in numpy-2
+// inside numpy/_core/include/numpy/numpyconfig.h
+#if NPY_API_VERSION >= 0x00000012 
+    #define USE_NUMPY_2
+#endif
+
 /**begin repeat
  * Float types
  *  #TYPE = FLOAT, DOUBLE#
@@ -271,9 +277,15 @@ MKL_UMATH_API
 void
 mkl_umath_@TYPE@_ldexp(char **args, const npy_intp *dimensions, const npy_intp *steps, void *NPY_UNUSED(func));
 
+#ifdef USE_NUMPY_2
+MKL_UMATH_API
+void
+mkl_umath_@TYPE@_ldexp_int64(char **args, const npy_intp *dimensions, const npy_intp *steps, void *NPY_UNUSED(func));
+#else
 MKL_UMATH_API
 void
 mkl_umath_@TYPE@_ldexp_long(char **args, const npy_intp *dimensions, const npy_intp *steps, void *NPY_UNUSED(func));
+#endif
 
 #define mkl_umath_@TYPE@_true_divide mkl_umath_@TYPE@_divide
 

mkl_umath/tests/test_basic.py

@@ -26,6 +26,7 @@
 import pytest
 import numpy as np
 import mkl_umath._ufuncs as mu
+import mkl_umath._patch as mp
 
 np.random.seed(42)
 
@@ -43,7 +44,9 @@ def get_args(args_str):
         elif s == 'i':
             args.append(np.int_(np.random.randint(low=1, high=10)))
         elif s == 'l':
-            args.append(np.dtype('long').type(np.random.randint(low=1, high=10)))
+            args.append(np.int64(np.random.randint(low=1, high=10)))
+        elif s == 'q':
+            args.append(np.int64(np.random.randint(low=1, high=10)))
         else:
             raise ValueError("Unexpected type specified!")
     return tuple(args)
@@ -82,6 +85,12 @@ def test_umath(case):
 
     assert np.allclose(mkl_res, np_res), f"Results for '{umath}': mkl_res: {mkl_res}, np_res: {np_res}"
 
-def test_cases_count():
-    print("Test cases count:", len(test_cases))
-    assert len(test_cases) > 0, "No test cases found"
+def test_patch():
+    mp.restore()
+    assert not mp.is_patched()
+
+    mp.use_in_numpy()  # Enable mkl_umath in Numpy
+    assert mp.is_patched()
+
+    mp.restore()  # Disable mkl_umath in Numpy
+    assert not mp.is_patched()