From 79ede33d679f7b00f5ecb34f2d0f0892a3bab892 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 9 Jun 2026 12:55:25 -0700 Subject: [PATCH 1/3] add getter for ufunc types NumPy 2.5+ changed how types property is exposed in NumPy, causing indexing on the types to break also fixes a bug where patching would attempt to free null pointers when patching process fails midway --- mkl_umath/src/_patch_numpy.pyx | 54 ++++++++++++++++++++++------------ 1 file changed, 36 insertions(+), 18 deletions(-) diff --git a/mkl_umath/src/_patch_numpy.pyx b/mkl_umath/src/_patch_numpy.pyx index 598a0167..27712bb4 100644 --- a/mkl_umath/src/_patch_numpy.pyx +++ b/mkl_umath/src/_patch_numpy.pyx @@ -40,6 +40,16 @@ from libc.stdlib cimport free, malloc cnp.import_umath() +cdef extern from *: + """ + #include "numpy/ufuncobject.h" + static inline char* _get_ufunc_types(PyObject *u) { + return (char *)((PyUFuncObject *)u)->types; + } + """ + char* _get_ufunc_types(object u) noexcept + + ctypedef struct function_info: cnp.PyUFuncGenericFunction original_function cnp.PyUFuncGenericFunction patch_function @@ -53,32 +63,39 @@ cdef class _patch_impl: functions_dict = dict() def __cinit__(self): - cdef int pi, oi + cdef int pi, oi, i, nargs + cdef int expected_count + cdef char* patch_types + cdef char* orig_types - umaths = [i for i in dir(mu) if isinstance(getattr(mu, i), np.ufunc)] + self.functions = NULL self.functions_count = 0 + + umaths = [x for x in dir(mu) if isinstance(getattr(mu, x), np.ufunc)] + expected_count = 0 for umath in umaths: mkl_umath_func = getattr(mu, umath) - self.functions_count += mkl_umath_func.ntypes + expected_count += mkl_umath_func.ntypes self.functions = malloc( - self.functions_count * sizeof(function_info) + expected_count * sizeof(function_info) ) - func_number = 0 for umath in umaths: patch_umath = getattr(mu, umath) c_patch_umath = patch_umath c_orig_umath = getattr(np, umath) nargs = c_patch_umath.nargs + patch_types = _get_ufunc_types(c_patch_umath) + orig_types = _get_ufunc_types(c_orig_umath) for pi in range(c_patch_umath.ntypes): oi = 0 while oi < c_orig_umath.ntypes: found = True - for i in range(c_patch_umath.nargs): + for i in range(nargs): if ( - c_patch_umath.types[pi * nargs + i] - != c_orig_umath.types[oi * nargs + i] + patch_types[pi * nargs + i] + != orig_types[oi * nargs + i] ): found = False break @@ -86,23 +103,23 @@ cdef class _patch_impl: break oi = oi + 1 if oi < c_orig_umath.ntypes: - self.functions[func_number].original_function = ( + self.functions[self.functions_count].original_function = ( c_orig_umath.functions[oi] ) - self.functions[func_number].patch_function = ( + self.functions[self.functions_count].patch_function = ( c_patch_umath.functions[pi] ) - self.functions[func_number].signature = ( + self.functions[self.functions_count].signature = ( malloc(nargs * sizeof(int)) ) for i in range(nargs): - self.functions[func_number].signature[i] = ( - c_patch_umath.types[pi * nargs + i] + self.functions[self.functions_count].signature[i] = ( + patch_types[pi * nargs + i] ) self.functions_dict[(umath, patch_umath.types[pi])] = ( - func_number + self.functions_count ) - func_number = func_number + 1 + self.functions_count += 1 else: raise RuntimeError( f"Unable to find original function for: {umath} " @@ -110,9 +127,10 @@ cdef class _patch_impl: ) def __dealloc__(self): - for i in range(self.functions_count): - free(self.functions[i].signature) - free(self.functions) + if self.functions is not NULL: + for i in range(self.functions_count): + free(self.functions[i].signature) + free(self.functions) cdef int _replace_loop( self, From df5762ea398d086f0daa799dd497c9c0a922b9a3 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Mon, 29 Jun 2026 08:27:42 -0700 Subject: [PATCH 2/3] add gh-226 to changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index a868bb7b..7b129cdd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed * Removed `numpy-base` dependency and `USE_NUMPY_BASE` environment variable from conda recipe [gh-200](https://github.com/IntelPython/mkl_umath/pull/200) +* Updated `mkl_umath` patching to work with changes to NumPy Cython API present in NumPy 2.5 [gh-226](https://github.com/IntelPython/mkl_umath/pull/226) ### Fixed From db4fe070fe11be02172cd498f641821662005631 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Mon, 29 Jun 2026 08:45:11 -0700 Subject: [PATCH 3/3] apply review comments --- mkl_umath/src/_patch_numpy.pyx | 23 +++++++++++++++++------ 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/mkl_umath/src/_patch_numpy.pyx b/mkl_umath/src/_patch_numpy.pyx index 27712bb4..4338bf2d 100644 --- a/mkl_umath/src/_patch_numpy.pyx +++ b/mkl_umath/src/_patch_numpy.pyx @@ -59,10 +59,10 @@ ctypedef struct function_info: cdef class _patch_impl: cdef int functions_count cdef function_info* functions - - functions_dict = dict() + cdef dict functions_dict def __cinit__(self): + self.functions_dict = {} cdef int pi, oi, i, nargs cdef int expected_count cdef char* patch_types @@ -77,15 +77,25 @@ cdef class _patch_impl: mkl_umath_func = getattr(mu, umath) expected_count += mkl_umath_func.ntypes - self.functions = malloc( - expected_count * sizeof(function_info) - ) + if expected_count > 0: + self.functions = malloc( + expected_count * sizeof(function_info) + ) + if self.functions is NULL: + raise MemoryError( + "Failed to allocate memory for function_info array" + ) for umath in umaths: patch_umath = getattr(mu, umath) c_patch_umath = patch_umath c_orig_umath = getattr(np, umath) + # nargs must be >=0 as no ufuncs have no arguments nargs = c_patch_umath.nargs + if nargs <= 0: + raise RuntimeError( + f"Invalid number of arguments for ufunc {umath}: {nargs}" + ) patch_types = _get_ufunc_types(c_patch_umath) orig_types = _get_ufunc_types(c_orig_umath) for pi in range(c_patch_umath.ntypes): @@ -129,7 +139,8 @@ cdef class _patch_impl: def __dealloc__(self): if self.functions is not NULL: for i in range(self.functions_count): - free(self.functions[i].signature) + if self.functions[i].signature is not NULL: + free(self.functions[i].signature) free(self.functions) cdef int _replace_loop(