Repository: numpy/numpy
Pull request: MAINT: simplify power fast path logic
This is an initial draft to resolve https://github.com/numpy/numpy/issues/27082.

I have removed the fast paths from array_power and am planning to implement them in the individual UFunc templates which seem to be few. This prevents the need for scalar extraction and reduces divergence from the UFunc. Currently, I have only implemented for integer loops, but would appreciate feedback before I proceed.

I hope I understood the issue correctly. Thank you!

Problem statement:
BUG: Power fast-paths have wrong and confusing promotion logic
The power fast-paths use checks like `scalar == 2` to call `square`.  Which was more correct before NEP 50 (but I am sure there were bugs also before), but less so after, if you pass in something like `np.int64(2)` and the array is, say `int32`.

That logic is a bit confusing and diverges it from the ufunc.

Hints from issue discussion:

Commit message: MAINT: Remove blocking regression test for power fast paths

Patch to review:
diff --git a/numpy/_core/src/multiarray/number.c b/numpy/_core/src/multiarray/number.c
index e6c04c1c9a9c..bec44ef4c4d6 100644
--- a/numpy/_core/src/multiarray/number.c
+++ b/numpy/_core/src/multiarray/number.c
@@ -328,167 +328,6 @@ array_inplace_matrix_multiply(PyArrayObject *self, PyObject *other)
     return res;
 }
 
-/*
- * Determine if object is a scalar and if so, convert the object
- * to a double and place it in the out_exponent argument
- * and return the "scalar kind" as a result.   If the object is
- * not a scalar (or if there are other error conditions)
- * return NPY_NOSCALAR, and out_exponent is undefined.
- */
-static NPY_SCALARKIND
-is_scalar_with_conversion(PyObject *o2, double* out_exponent)
-{
-    PyObject *temp;
-    const int optimize_fpexps = 1;
-
-    if (PyLong_Check(o2)) {
-        long tmp = PyLong_AsLong(o2);
-        if (error_converting(tmp)) {
-            PyErr_Clear();
-            return NPY_NOSCALAR;
-        }
-        *out_exponent = (double)tmp;
-        return NPY_INTPOS_SCALAR;
-    }
-
-    if (optimize_fpexps && PyFloat_Check(o2)) {
-        *out_exponent = PyFloat_AsDouble(o2);
-        return NPY_FLOAT_SCALAR;
-    }
-
-    if (PyArray_Check(o2)) {
-        if ((PyArray_NDIM((PyArrayObject *)o2) == 0) &&
-                ((PyArray_ISINTEGER((PyArrayObject *)o2) ||
-                 (optimize_fpexps && PyArray_ISFLOAT((PyArrayObject *)o2))))) {
-            temp = Py_TYPE(o2)->tp_as_number->nb_float(o2);
-            if (temp == NULL) {
-                return NPY_NOSCALAR;
-            }
-            *out_exponent = PyFloat_AsDouble(o2);
-            Py_DECREF(temp);
-            if (PyArray_ISINTEGER((PyArrayObject *)o2)) {
-                return NPY_INTPOS_SCALAR;
-            }
-            else { /* ISFLOAT */
-                return NPY_FLOAT_SCALAR;
-            }
-        }
-    }
-    else if (PyArray_IsScalar(o2, Integer) ||
-                (optimize_fpexps && PyArray_IsScalar(o2, Floating))) {
-        temp = Py_TYPE(o2)->tp_as_number->nb_float(o2);
-        if (temp == NULL) {
-            return NPY_NOSCALAR;
-        }
-        *out_exponent = PyFloat_AsDouble(o2);
-        Py_DECREF(temp);
-
-        if (PyArray_IsScalar(o2, Integer)) {
-                return NPY_INTPOS_SCALAR;
-        }
-        else { /* IsScalar(o2, Floating) */
-            return NPY_FLOAT_SCALAR;
-        }
-    }
-    else if (PyIndex_Check(o2)) {
-        PyObject* value = PyNumber_Index(o2);
-        Py_ssize_t val;
-        if (value == NULL) {
-            if (PyErr_Occurred()) {
-                PyErr_Clear();
-            }
-            return NPY_NOSCALAR;
-        }
-        val = PyLong_AsSsize_t(value);
-        Py_DECREF(value);
-        if (error_converting(val)) {
-            PyErr_Clear();
-            return NPY_NOSCALAR;
-        }
-        *out_exponent = (double) val;
-        return NPY_INTPOS_SCALAR;
-    }
-    return NPY_NOSCALAR;
-}
-
-/*
- * optimize float array or complex array to a scalar power
- * returns 0 on success, -1 if no optimization is possible
- * the result is in value (can be NULL if an error occurred)
- */
-static int
-fast_scalar_power(PyObject *o1, PyObject *o2, int inplace,
-                  PyObject **value)
-{
-    double exponent;
-    NPY_SCALARKIND kind;   /* NPY_NOSCALAR is not scalar */
-
-    if (PyArray_Check(o1) &&
-            !PyArray_ISOBJECT((PyArrayObject *)o1) &&
-            ((kind=is_scalar_with_conversion(o2, &exponent))>0)) {
-        PyArrayObject *a1 = (PyArrayObject *)o1;
-        PyObject *fastop = NULL;
-        if (PyArray_ISFLOAT(a1) || PyArray_ISCOMPLEX(a1)) {
-            if (exponent == 1.0) {
-                fastop = n_ops.positive;
-            }
-            else if (exponent == -1.0) {
-                fastop = n_ops.reciprocal;
-            }
-            else if (exponent ==  0.0) {
-                fastop = n_ops._ones_like;
-            }
-            else if (exponent ==  0.5) {
-                fastop = n_ops.sqrt;
-            }
-            else if (exponent ==  2.0) {
-                fastop = n_ops.square;
-            }
-            else {
-                return -1;
-            }
-
-            if (inplace || can_elide_temp_unary(a1)) {
-                *value = PyArray_GenericInplaceUnaryFunction(a1, fastop);
-            }
-            else {
-                *value = PyArray_GenericUnaryFunction(a1, fastop);
-            }
-            return 0;
-        }
-        /* Because this is called with all arrays, we need to
-         *  change the output if the kind of the scalar is different
-         *  than that of the input and inplace is not on ---
-         *  (thus, the input should be up-cast)
-         */
-        else if (exponent == 2.0) {
-            fastop = n_ops.square;
-            if (inplace) {
-                *value = PyArray_GenericInplaceUnaryFunction(a1, fastop);
-            }
-            else {
-                /* We only special-case the FLOAT_SCALAR and integer types */
-                if (kind == NPY_FLOAT_SCALAR && PyArray_ISINTEGER(a1)) {
-                    PyArray_Descr *dtype = PyArray_DescrFromType(NPY_DOUBLE);
-                    a1 = (PyArrayObject *)PyArray_CastToType(a1, dtype,
-                            PyArray_ISFORTRAN(a1));
-                    if (a1 != NULL) {
-                        /* cast always creates a new array */
-                        *value = PyArray_GenericInplaceUnaryFunction(a1, fastop);
-                        Py_DECREF(a1);
-                    }
-                }
-                else {
-                    *value = PyArray_GenericUnaryFunction(a1, fastop);
-                }
-            }
-            return 0;
-        }
-    }
-    /* no fast operation found */
-    return -1;
-}
-
 static PyObject *
 array_power(PyObject *a1, PyObject *o2, PyObject *modulo)
 {
@@ -501,9 +340,7 @@ array_power(PyObject *a1, PyObject *o2, PyObject *modulo)
     }
 
     BINOP_GIVE_UP_IF_NEEDED(a1, o2, nb_power, array_power);
-    if (fast_scalar_power(a1, o2, 0, &value) != 0) {
-        value = PyArray_GenericBinaryFunction(a1, o2, n_ops.power);
-    }
+    value = PyArray_GenericBinaryFunction(a1, o2, n_ops.power);
     return value;
 }
 
@@ -643,9 +480,7 @@ array_inplace_power(PyArrayObject *a1, PyObject *o2, PyObject *NPY_UNUSED(modulo
 
     INPLACE_GIVE_UP_IF_NEEDED(
             a1, o2, nb_inplace_power, array_inplace_power);
-    if (fast_scalar_power((PyObject *)a1, o2, 1, &value) != 0) {
-        value = PyArray_GenericInplaceBinaryFunction(a1, o2, n_ops.power);
-    }
+    value = PyArray_GenericInplaceBinaryFunction(a1, o2, n_ops.power);
     return value;
 }
 
diff --git a/numpy/_core/src/umath/loops.c.src b/numpy/_core/src/umath/loops.c.src
index 5ac67fa3024b..cca60062b361 100644
--- a/numpy/_core/src/umath/loops.c.src
+++ b/numpy/_core/src/umath/loops.c.src
@@ -524,16 +524,24 @@ NPY_NO_EXPORT void
 #endif
         if (in2 == 0) {
             *((@type@ *)op1) = 1;
-            continue;
+        } 
+        else if (in2 == 1)
+        {
+            *((@type@ *)op1) = in1;
+        }
+        else if (in2 == 2)
+        {
+            *((@type@ *)op1) = in1 * in1;
         }
-        if (in1 == 1) {
+        else if (in1 == 1) {
             *((@type@ *)op1) = 1;
-            continue;
         }
-
-        int first_bit = in2 & 1;
-        in2 >>= 1;
-        *((@type@ *) op1) = _@TYPE@_squared_exponentiation_helper(in1, in2, first_bit);
+        else
+        {
+            int first_bit = in2 & 1;
+            in2 >>= 1;
+            *((@type@ *) op1) = _@TYPE@_squared_exponentiation_helper(in1, in2, first_bit);
+        }
     }
 }
 /**end repeat**/
diff --git a/numpy/_core/tests/test_multiarray.py b/numpy/_core/tests/test_multiarray.py
index 02ed3ece94b5..b17e80792753 100644
--- a/numpy/_core/tests/test_multiarray.py
+++ b/numpy/_core/tests/test_multiarray.py
@@ -4123,27 +4123,6 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kw):
         assert_equal(A[0], 30)
         assert_(isinstance(A, OutClass))
 
-    def test_pow_override_with_errors(self):
-        # regression test for gh-9112
-        class PowerOnly(np.ndarray):
-            def __array_ufunc__(self, ufunc, method, *inputs, **kw):
-                if ufunc is not np.power:
-                    raise NotImplementedError
-                return "POWER!"
-        # explicit cast to float, to ensure the fast power path is taken.
-        a = np.array(5., dtype=np.float64).view(PowerOnly)
-        assert_equal(a ** 2.5, "POWER!")
-        with assert_raises(NotImplementedError):
-            a ** 0.5
-        with assert_raises(NotImplementedError):
-            a ** 0
-        with assert_raises(NotImplementedError):
-            a ** 1
-        with assert_raises(NotImplementedError):
-            a ** -1
-        with assert_raises(NotImplementedError):
-            a ** 2
-
     def test_pow_array_object_dtype(self):
         # test pow on arrays of object dtype
         class SomeClass:


Code review comments:
diff_hunk: @@ -524,16 +524,24 @@ NPY_NO_EXPORT void
 #endif
         if (in2 == 0) {
             *((@type@ *)op1) = 1;
-            continue;
+        } 
+        else if (in2 == 1)
+        {
+            *((@type@ *)op1) = in1;
+        }
+        else if (in2 == 2)
+        {
+            *((@type@ *)op1) = in1 * in1;
         }
-        if (in1 == 1) {
+        else if (in1 == 1) {
             *((@type@ *)op1) = 1;
path: numpy/_core/src/umath/loops.c.src
text: @user1:
I am a bit confused by the new (and existing) fast paths here.  The first branch is the scalar case and it would seem to me that any fast path is even more relevant there?
Should the fast path just be copied?  Maybe there should be a second helper with the body so that we can have a `if (stride[1] == 0) {call helper()} else {call_helper()}` to nudge the compiler to optimize for 0 stride (and assume it'll lift the checks out of the loop then).

@user1:
To be clear, the `if (steps[1]==0) {` path (sorry, I wrote `stride` above, the ufunc code here calls it "steps") is the path we need to worry about being fast.

I.e. if the second operand is a scalar, then we will always take that part (I am not 100% sure about `scalar**scalar`.).

@author:
Yes, I stuck to the original fast path logic but I agree the scalar case is more important for efficiency. I've created a helper function to remove repetition.