TST: Add basic tests for custom DType (scaled float) ufuncs

This does not actually test a lot of new paths, since the normal
ufuncs do cover most of the typical paths.

One thing it does test is that the ufunc implementation (as in
ArrayMethod) can customize the cast-safety.  This is used here
since the parametric DType considers different scalings as same-kind
casts and not safe/equiv casts.

Author: seberg

numpy/core/src/multiarray/array_method.c

@@ -833,10 +833,12 @@ generic_masked_strided_loop(PyArrayMethod_Context *context,
 
 
 /*
- * Identical to the `get_loop` functions and wraps it.  This adds support
- * to a boolean mask being passed in as a last, additional, operand.
- * The wrapped loop will only be called for unmasked elements.
- * (Does not support `move_references` or inner dimensions!)
+ * Fetches a strided-loop function that supports a boolean mask as additional
+ * (last) operand to the strided-loop.  It is otherwise largely identical to
+ * the `get_loop` method which it wraps.
+ * This is the core implementation for the ufunc `where=...` keyword argument.
+ *
+ * NOTE: This function does not support `move_references` or inner dimensions.
  */
 NPY_NO_EXPORT int
 PyArrayMethod_GetMaskedStridedLoop(

numpy/core/src/umath/_scaled_float_dtype.c

@@ -23,6 +23,7 @@
 #include "numpy/npy_math.h"
 #include "convert_datatype.h"
 #include "dtypemeta.h"
+#include "dispatching.h"
 
 
 typedef struct {
@@ -456,6 +457,225 @@ init_casts(void)
 }
 
 
+/*
+ * We also wish to test very simple ufunc functionality.  So create two
+ * ufunc loops:
+ * 1. Multiplication, which can multiply the factors and work with that.
+ * 2. Addition, which needs to use the common instance, and runs into
+ *    cast safety subtleties since we will implement it without an additional
+ *    cast.
+ *
+ * NOTE: When first writing this, promotion did not exist for new-style loops,
+ *       if it exists, we could use promotion to implement double * sfloat.
+ */
+static int
+multiply_sfloats(PyArrayMethod_Context *NPY_UNUSED(context),
+        char *const data[], npy_intp const dimensions[],
+        npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
+{
+    npy_intp N = dimensions[0];
+    char *in1 = data[0];
+    char *in2 = data[1];
+    char *out = data[2];
+    for (npy_intp i = 0; i < N; i++) {
+        *(double *)out = *(double *)in1 * *(double *)in2;
+        in1 += strides[0];
+        in2 += strides[1];
+        out += strides[2];
+    }
+    return 0;
+}
+
+
+static NPY_CASTING
+multiply_sfloats_resolve_descriptors(
+        PyArrayMethodObject *NPY_UNUSED(self),
+        PyArray_DTypeMeta *NPY_UNUSED(dtypes[3]),
+        PyArray_Descr *given_descrs[3],
+        PyArray_Descr *loop_descrs[3])
+{
+    /*
+     * Multiply the scaling for the result.  If the result was passed in we
+     * simply ignore it and let the casting machinery fix it up here.
+     */
+    double factor = ((PyArray_SFloatDescr *)given_descrs[1])->scaling;
+    loop_descrs[2] = sfloat_scaled_copy(
+            (PyArray_SFloatDescr *)given_descrs[0], factor);
+    if (loop_descrs[2] == 0) {
+        return -1;
+    }
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[0] = given_descrs[0];
+    Py_INCREF(given_descrs[1]);
+    loop_descrs[1] = given_descrs[1];
+    return NPY_NO_CASTING;
+}
+
+
+/*
+ * Unlike the multiplication implementation above, this loops deals with
+ * scaling (casting) internally.  This allows to test some different paths.
+ */
+static int
+add_sfloats(PyArrayMethod_Context *context,
+        char *const data[], npy_intp const dimensions[],
+        npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
+{
+    double fin1 = ((PyArray_SFloatDescr *)context->descriptors[0])->scaling;
+    double fin2 = ((PyArray_SFloatDescr *)context->descriptors[1])->scaling;
+    double fout = ((PyArray_SFloatDescr *)context->descriptors[2])->scaling;
+
+    double fact1 = fin1 / fout;
+    double fact2 = fin2 / fout;
+    if (check_factor(fact1) < 0) {
+        return -1;
+    }
+    if (check_factor(fact2) < 0) {
+        return -1;
+    }
+
+    npy_intp N = dimensions[0];
+    char *in1 = data[0];
+    char *in2 = data[1];
+    char *out = data[2];
+    for (npy_intp i = 0; i < N; i++) {
+        *(double *)out = (*(double *)in1 * fact1) + (*(double *)in2 * fact2);
+        in1 += strides[0];
+        in2 += strides[1];
+        out += strides[2];
+    }
+    return 0;
+}
+
+
+static NPY_CASTING
+add_sfloats_resolve_descriptors(
+        PyArrayMethodObject *NPY_UNUSED(self),
+        PyArray_DTypeMeta *NPY_UNUSED(dtypes[3]),
+        PyArray_Descr *given_descrs[3],
+        PyArray_Descr *loop_descrs[3])
+{
+    /*
+     * Here we accept an output descriptor (the inner loop can deal with it),
+     * if none is given, we use the "common instance":
+     */
+    if (given_descrs[2] == NULL) {
+        loop_descrs[2] = sfloat_common_instance(
+                given_descrs[0], given_descrs[1]);
+        if (loop_descrs[2] == 0) {
+            return -1;
+        }
+    }
+    else {
+        Py_INCREF(given_descrs[2]);
+        loop_descrs[2] = given_descrs[2];
+    }
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[0] = given_descrs[0];
+    Py_INCREF(given_descrs[1]);
+    loop_descrs[1] = given_descrs[1];
+
+    /* If the factors mismatch, we do implicit casting inside the ufunc! */
+    double fin1 = ((PyArray_SFloatDescr *)loop_descrs[0])->scaling;
+    double fin2 = ((PyArray_SFloatDescr *)loop_descrs[1])->scaling;
+    double fout = ((PyArray_SFloatDescr *)loop_descrs[2])->scaling;
+
+    if (fin1 == fout && fin2 == fout) {
+        return NPY_NO_CASTING;
+    }
+    if (npy_fabs(fin1) == npy_fabs(fout) && npy_fabs(fin2) == npy_fabs(fout)) {
+        return NPY_EQUIV_CASTING;
+    }
+    return NPY_SAME_KIND_CASTING;
+}
+
+
+static int
+add_loop(const char *ufunc_name, PyBoundArrayMethodObject *bmeth)
+{
+    PyObject *mod = PyImport_ImportModule("numpy");
+    if (mod == NULL) {
+        return -1;
+    }
+    PyObject *ufunc = PyObject_GetAttrString(mod, ufunc_name);
+    Py_DECREF(mod);
+    if (!PyObject_TypeCheck(ufunc, &PyUFunc_Type)) {
+        Py_DECREF(ufunc);
+        PyErr_Format(PyExc_TypeError,
+                "numpy.%s was not a ufunc!", ufunc_name);
+        return -1;
+    }
+    PyObject *dtype_tup = PyArray_TupleFromItems(
+            3, (PyObject **)bmeth->dtypes, 0);
+    if (dtype_tup == NULL) {
+        Py_DECREF(ufunc);
+        return -1;
+    }
+    PyObject *info = PyTuple_Pack(2, dtype_tup, bmeth->method);
+    Py_DECREF(dtype_tup);
+    if (info == NULL) {
+        Py_DECREF(ufunc);
+        return -1;
+    }
+    int res = PyUFunc_AddLoop((PyUFuncObject *)ufunc, info, 0);
+    Py_DECREF(ufunc);
+    Py_DECREF(info);
+    return res;
+}
+
+
+/*
+ * Add new ufunc loops (this is somewhat clumsy as of writing it, but should
+ * get less so with the introduction of public API).
+ */
+static int
+init_ufuncs(void) {
+    PyArray_DTypeMeta *dtypes[3] = {
+            &PyArray_SFloatDType, &PyArray_SFloatDType, &PyArray_SFloatDType};
+    PyType_Slot slots[3] = {{0, NULL}};
+    PyArrayMethod_Spec spec = {
+        .nin = 2,
+        .nout =1,
+        .dtypes = dtypes,
+        .slots = slots,
+    };
+    spec.name = "sfloat_multiply";
+    spec.casting = NPY_NO_CASTING;
+
+    slots[0].slot = NPY_METH_resolve_descriptors;
+    slots[0].pfunc = &multiply_sfloats_resolve_descriptors;
+    slots[1].slot = NPY_METH_strided_loop;
+    slots[1].pfunc = &multiply_sfloats;
+    PyBoundArrayMethodObject *bmeth = PyArrayMethod_FromSpec_int(&spec, 0);
+    if (bmeth == NULL) {
+        return -1;
+    }
+    int res = add_loop("multiply", bmeth);
+    Py_DECREF(bmeth);
+    if (res < 0) {
+        return -1;
+    }
+
+    spec.name = "sfloat_add";
+    spec.casting = NPY_SAME_KIND_CASTING;
+
+    slots[0].slot = NPY_METH_resolve_descriptors;
+    slots[0].pfunc = &add_sfloats_resolve_descriptors;
+    slots[1].slot = NPY_METH_strided_loop;
+    slots[1].pfunc = &add_sfloats;
+    bmeth = PyArrayMethod_FromSpec_int(&spec, 0);
+    if (bmeth == NULL) {
+        return -1;
+    }
+    res = add_loop("add", bmeth);
+    Py_DECREF(bmeth);
+    if (res < 0) {
+        return -1;
+    }
+    return 0;
+}
+
+
 /*
  * Python entry point, exported via `umathmodule.h` and `multiarraymodule.c`.
  * TODO: Should be moved when the necessary API is not internal anymore.
@@ -491,6 +711,10 @@ get_sfloat_dtype(PyObject *NPY_UNUSED(mod), PyObject *NPY_UNUSED(args))
         return NULL;
     }
 
+    if (init_ufuncs() < 0) {
+        return NULL;
+    }
+
     initalized = NPY_TRUE;
     return (PyObject *)&PyArray_SFloatDType;
 }

numpy/core/src/umath/dispatching.c

@@ -69,8 +69,8 @@ promote_and_get_info_and_ufuncimpl(PyUFuncObject *ufunc,
  * @param ignore_duplicate If 1 and a loop with the same `dtype_tuple` is
  *        found, the function does nothing.
  */
-static int
-add_ufunc_loop(PyUFuncObject *ufunc, PyObject *info, int ignore_duplicate)
+NPY_NO_EXPORT int
+PyUFunc_AddLoop(PyUFuncObject *ufunc, PyObject *info, int ignore_duplicate)
 {
     /*
      * Validate the info object, this should likely move to to a different
@@ -495,7 +495,7 @@ add_and_return_legacy_wrapping_ufunc_loop(PyUFuncObject *ufunc,
     if (info == NULL) {
         return NULL;
     }
-    if (add_ufunc_loop(ufunc, info, ignore_duplicate) < 0) {
+    if (PyUFunc_AddLoop(ufunc, info, ignore_duplicate) < 0) {
         Py_DECREF(info);
         return NULL;
     }

numpy/core/src/umath/dispatching.h

@@ -7,6 +7,9 @@
 #include "array_method.h"
 
 
+NPY_NO_EXPORT int
+PyUFunc_AddLoop(PyUFuncObject *ufunc, PyObject *info, int ignore_duplicate);
+
 NPY_NO_EXPORT PyArrayMethodObject *
 promote_and_get_ufuncimpl(PyUFuncObject *ufunc,
         PyArrayObject *const ops[],

numpy/core/tests/test_custom_dtypes.py

@@ -90,3 +90,46 @@ def test_sfloat_promotion(self):
         # Test an undefined promotion:
         with pytest.raises(TypeError):
             np.result_type(SF(1.), np.int64)
+
+    def test_basic_multiply(self):
+        a = self._get_array(2.)
+        b = self._get_array(4.)
+
+        res = a * b
+        # multiplies dtype scaling and content separately:
+        assert res.dtype.get_scaling() == 8.
+        expected_view = a.view(np.float64) * b.view(np.float64)
+        assert_array_equal(res.view(np.float64), expected_view)
+
+    def test_basic_addition(self):
+        a = self._get_array(2.)
+        b = self._get_array(4.)
+
+        res = a + b
+        # addition uses the type promotion rules for the result:
+        assert res.dtype == np.result_type(a.dtype, b.dtype)
+        expected_view = (a.astype(res.dtype).view(np.float64) +
+                         b.astype(res.dtype).view(np.float64))
+        assert_array_equal(res.view(np.float64), expected_view)
+
+    def test_addition_cast_safety(self):
+        """The addition method is special for the scaled float, because it
+        includes the "cast" between different factors, thus cast-safety
+        is influenced by the implementation.
+        """
+        a = self._get_array(2.)
+        b = self._get_array(-2.)
+        c = self._get_array(3.)
+
+        # sign change is "equiv":
+        np.add(a, b, casting="equiv")
+        with pytest.raises(TypeError):
+            np.add(a, b, casting="no")
+
+        # Different factor is "same_kind" (default) so check that "safe" fails
+        with pytest.raises(TypeError):
+            np.add(a, c, casting="safe")
+
+        # Check that casting the output fails also (done by the ufunc here)
+        with pytest.raises(TypeError):
+            np.add(a, a, out=c, casting="safe")