numpy · ngoldbaum · Oct 19, 2023 · Oct 10, 2023 · Oct 12, 2023 · Oct 12, 2023
diff --git a/numpy/_core/include/numpy/_dtype_api.h b/numpy/_core/include/numpy/_dtype_api.h
@@ -5,7 +5,7 @@
 #ifndef NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_
 #define NUMPY_CORE_INCLUDE_NUMPY___DTYPE_API_H_
 
-#define __EXPERIMENTAL_DTYPE_API_VERSION 13
+#define __EXPERIMENTAL_DTYPE_API_VERSION 14
 
 struct PyArrayMethodObject_tag;
 
@@ -129,16 +129,17 @@ typedef struct {
  * SLOTS IDs For the ArrayMethod creation, once fully public, IDs are fixed
  * but can be deprecated and arbitrarily extended.
  */
-#define NPY_METH_resolve_descriptors 1
+#define _NPY_METH_resolve_descriptors_with_scalars 1
+#define NPY_METH_resolve_descriptors 2
 /* We may want to adapt the `get_loop` signature a bit: */
-#define _NPY_METH_get_loop 2
-#define NPY_METH_get_reduction_initial 3
+#define _NPY_METH_get_loop 3
+#define NPY_METH_get_reduction_initial 4
 /* specific loops for constructions/default get_loop: */
-#define NPY_METH_strided_loop 4
-#define NPY_METH_contiguous_loop 5
-#define NPY_METH_unaligned_strided_loop 6
-#define NPY_METH_unaligned_contiguous_loop 7
-#define NPY_METH_contiguous_indexed_loop 8
+#define NPY_METH_strided_loop 5
+#define NPY_METH_contiguous_loop 6
+#define NPY_METH_unaligned_strided_loop 7
+#define NPY_METH_unaligned_contiguous_loop 8
+#define NPY_METH_contiguous_indexed_loop 9
 
 /*
  * The resolve descriptors function, must be able to handle NULL values for
@@ -162,6 +163,30 @@ typedef NPY_CASTING (resolve_descriptors_function)(
         npy_intp *view_offset);
 
 
+/*
+ * Rarely needed, slightly more powerful version of `resolve_descriptors`.
+ * See also `resolve_descriptors_function` for details on shared arguments.
+ *
+ * NOTE: This function is private now as it is unclear how and what to pass
+ *       exactly as additional information to allow dealing with the scalars.
+ *       See also gh-24915.
+ */
+typedef NPY_CASTING (resolve_descriptors_with_scalars_function)(
+        struct PyArrayMethodObject_tag *method,
+        PyArray_DTypeMeta **dtypes,
+        /* Unlike above, these can have any DType and we may allow NULL. */
+        PyArray_Descr **given_descrs,
+        /*
+         * Input scalars or NULL.  Only ever passed for python scalars.
+         * WARNING: In some cases, a loop may be explicitly selected and the
+         *     value passed is not available (NULL) or does not have the
+         *     expected type.
+         */
+        PyObject *const *input_scalars,
+        PyArray_Descr **loop_descrs,
+        npy_intp *view_offset);
+
+
 typedef int (PyArrayMethod_StridedLoop)(PyArrayMethod_Context *context,
         char *const *data, const npy_intp *dimensions, const npy_intp *strides,
         NpyAuxData *transferdata);

diff --git a/numpy/_core/meson.build b/numpy/_core/meson.build
@@ -1097,6 +1097,7 @@ src_umath = umath_gen_headers + [
   src_file.process('src/umath/scalarmath.c.src'),
   'src/umath/ufunc_object.c',
   'src/umath/umathmodule.c',
+  'src/umath/special_integer_comparisons.cpp',
   'src/umath/string_ufuncs.cpp',
   'src/umath/wrapping_array_method.c',
   # For testing. Eventually, should use public API and be separate:

diff --git a/numpy/_core/src/multiarray/abstractdtypes.h b/numpy/_core/src/multiarray/abstractdtypes.h
@@ -5,6 +5,10 @@
 #include "dtypemeta.h"
 
 
+#ifdef __cplusplus
+extern "C" {
+#endif
+
 /*
  * These are mainly needed for value based promotion in ufuncs.  It
  * may be necessary to make them (partially) public, to allow user-defined
@@ -70,4 +74,8 @@ npy_mark_tmp_array_if_pyscalar(
     return 0;
 }
 
+#ifdef __cplusplus
+}
+#endif
+
 #endif  /* NUMPY_CORE_SRC_MULTIARRAY_ABSTRACTDTYPES_H_ */
diff --git a/numpy/_core/src/multiarray/array_method.c b/numpy/_core/src/multiarray/array_method.c
@@ -221,12 +221,6 @@ validate_spec(PyArrayMethod_Spec *spec)
                     "(method: %s)", spec->dtypes[i], spec->name);
             return -1;
         }
-        if (NPY_DT_is_abstract(spec->dtypes[i])) {
-            PyErr_Format(PyExc_TypeError,
-                    "abstract DType %S are currently not supported."
-                    "(method: %s)", spec->dtypes[i], spec->name);
-            return -1;
-        }
     }
     return 0;
 }
@@ -261,6 +255,16 @@ fill_arraymethod_from_slots(
      */
     for (PyType_Slot *slot = &spec->slots[0]; slot->slot != 0; slot++) {
         switch (slot->slot) {
+            case _NPY_METH_resolve_descriptors_with_scalars:
+                if (!private) {
+                    PyErr_SetString(PyExc_RuntimeError,
+                            "the _NPY_METH_resolve_descriptors_with_scalars "
+                            "slot private due to uncertainty about the best "
+                            "signature (see gh-24915)");
+                    return -1;
+                }
+                meth->resolve_descriptors_with_scalars = slot->pfunc;
+                continue;
             case NPY_METH_resolve_descriptors:
                 meth->resolve_descriptors = slot->pfunc;
                 continue;
@@ -272,7 +276,6 @@ fill_arraymethod_from_slots(
                  *       (as in: we should not worry about changing it, but of
                  *       course that would not break it immediately.)
                  */
-                /* Only allow override for private functions initially */
                 meth->get_strided_loop = slot->pfunc;
                 continue;
             /* "Typical" loops, supported used by the default `get_loop` */

diff --git a/numpy/_core/src/multiarray/array_method.h b/numpy/_core/src/multiarray/array_method.h
@@ -45,6 +45,7 @@ typedef struct PyArrayMethodObject_tag {
     NPY_CASTING casting;
     /* default flags. The get_strided_loop function can override these */
     NPY_ARRAYMETHOD_FLAGS flags;
+    resolve_descriptors_with_scalars_function *resolve_descriptors_with_scalars;
     resolve_descriptors_function *resolve_descriptors;
     get_loop_function *get_strided_loop;
     get_reduction_initial_function  *get_reduction_initial;

diff --git a/numpy/_core/src/multiarray/arrayobject.h b/numpy/_core/src/multiarray/arrayobject.h
@@ -5,6 +5,10 @@
 #ifndef NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_
 #define NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_
 
+#ifdef __cplusplus
+extern "C" {
+#endif
+
 extern NPY_NO_EXPORT npy_bool numpy_warn_if_no_mem_policy;
 
 NPY_NO_EXPORT PyObject *
@@ -51,4 +55,8 @@ static const int NPY_ARRAY_WAS_PYTHON_COMPLEX = (1 << 28);
 static const int NPY_ARRAY_WAS_INT_AND_REPLACED = (1 << 27);
 static const int NPY_ARRAY_WAS_PYTHON_LITERAL = (1 << 30 | 1 << 29 | 1 << 28);
 
+#ifdef __cplusplus
+}
+#endif
+
 #endif  /* NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_ */
diff --git a/numpy/_core/src/umath/dispatching.c b/numpy/_core/src/umath/dispatching.c
@@ -152,14 +152,21 @@ PyUFunc_AddLoop(PyUFuncObject *ufunc, PyObject *info, int ignore_duplicate)
  */
 NPY_NO_EXPORT int
 PyUFunc_AddLoopFromSpec(PyObject *ufunc, PyArrayMethod_Spec *spec)
+{
+    return PyUFunc_AddLoopFromSpec_int(ufunc, spec, 0);
+}
+
+
+NPY_NO_EXPORT int
+PyUFunc_AddLoopFromSpec_int(PyObject *ufunc, PyArrayMethod_Spec *spec, int priv)
 {
     if (!PyObject_TypeCheck(ufunc, &PyUFunc_Type)) {
         PyErr_SetString(PyExc_TypeError,
                 "ufunc object passed is not a ufunc!");
         return -1;
     }
     PyBoundArrayMethodObject *bmeth =
-            (PyBoundArrayMethodObject *)PyArrayMethod_FromSpec(spec);
+            (PyBoundArrayMethodObject *)PyArrayMethod_FromSpec_int(spec, priv);
     if (bmeth == NULL) {
         return -1;
     }
@@ -275,6 +282,12 @@ resolve_implementation_info(PyUFuncObject *ufunc,
                         == PyTuple_GET_ITEM(curr_dtypes, 2)) {
                     continue;
                 }
+                /*
+                 * This should be a reduce, but doesn't follow the reduce
+                 * pattern.  So (for now?) consider this not a match.
+                 */
+                matches = NPY_FALSE;
+                continue;
             }
 
             if (resolver_dtype == (PyArray_DTypeMeta *)Py_None) {

diff --git a/numpy/_core/src/umath/dispatching.h b/numpy/_core/src/umath/dispatching.h
@@ -20,6 +20,9 @@ PyUFunc_AddLoop(PyUFuncObject *ufunc, PyObject *info, int ignore_duplicate);
 NPY_NO_EXPORT int
 PyUFunc_AddLoopFromSpec(PyObject *ufunc, PyArrayMethod_Spec *spec);
 
+NPY_NO_EXPORT int
+PyUFunc_AddLoopFromSpec_int(PyObject *ufunc, PyArrayMethod_Spec *spec, int priv);
+
 NPY_NO_EXPORT PyArrayMethodObject *
 promote_and_get_ufuncimpl(PyUFuncObject *ufunc,
         PyArrayObject *const ops[],

diff --git a/numpy/_core/src/umath/legacy_array_method.h b/numpy/_core/src/umath/legacy_array_method.h
@@ -5,13 +5,14 @@
 #include "numpy/ufuncobject.h"
 #include "array_method.h"
 
+#ifdef __cplusplus
+extern "C" {
+#endif
 
 NPY_NO_EXPORT PyArrayMethodObject *
 PyArray_NewLegacyWrappingArrayMethod(PyUFuncObject *ufunc,
         PyArray_DTypeMeta *signature[]);
 
-
-
 /*
  * The following two symbols are in the header so that other places can use
  * them to probe for special cases (or whether an ArrayMethod is a "legacy"
@@ -29,5 +30,8 @@ NPY_NO_EXPORT NPY_CASTING
 wrapped_legacy_resolve_descriptors(PyArrayMethodObject *,
         PyArray_DTypeMeta **, PyArray_Descr **, PyArray_Descr **, npy_intp *);
 
+#ifdef __cplusplus
+}
+#endif
 
 #endif  /*_NPY_LEGACY_ARRAY_METHOD_H */
diff --git a/numpy/_core/src/umath/scalarmath.c.src b/numpy/_core/src/umath/scalarmath.c.src
@@ -1842,6 +1842,7 @@ static PyObject *
  *         LONG, ULONG, LONGLONG, ULONGLONG,
  *         HALF, FLOAT, DOUBLE, LONGDOUBLE,
  *         CFLOAT, CDOUBLE, CLONGDOUBLE#
+ * #isint = 1*10, 0*7#
  * #simp = def*10, def_half, def*3, fcmplx, cmplx, lcmplx#
  */
 #define IS_@name@
@@ -1852,6 +1853,27 @@ static PyObject*
     npy_@name@ arg1, arg2;
     int out = 0;
 
+#if @isint@
+    /* Special case comparison with python integers */
+    if (PyLong_CheckExact(other)) {
+        PyObject *self_val = PyNumber_Index(self);
+        if (self_val == NULL) {
+            return NULL;
+        }
+        int res = PyObject_RichCompareBool(self_val, other, cmp_op);
+        Py_DECREF(self_val);
+        if (res < 0) {
+            return NULL;
+        }
+        else if (res) {
+            PyArrayScalar_RETURN_TRUE;
+        }
+        else {
+            PyArrayScalar_RETURN_FALSE;
+        }
+    }
+#endif
+
     /*
      * Extract the other value (if it is compatible).
      */