ENH: added check for scalar array in get_ufunc_arguments from ufunc_object.c

Check if object is of array dimension zero, then avoid heavy calls like PyArray_FromAny. Hence passing obj to out_op PyArrayObject unchanged. Also PyArray_FromAny(obj,NULL,0,0,0,NULL) calls PyArray_FromArray(obj,NULL,0) which return original obj

1 files changed, 25 insertions, 14 deletions

diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index 85c0069c3..1d003842c 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -753,22 +753,33 @@ static int get_ufunc_arguments(PyUFuncObject *ufunc,
     /* Get input arguments */
     for (i = 0; i < nin; ++i) {
         obj = PyTuple_GET_ITEM(args, i);
-        if (!PyArray_Check(obj) && !PyArray_IsScalar(obj, Generic)) {
-            /*
-             * TODO: There should be a comment here explaining what
-             *       context does.
-             */
-            context = Py_BuildValue("OOi", ufunc, args, i);
-            if (context == NULL) {
-                return -1;
+ 
+        /*
+         * If obj is array of dimension zero. PyArray_FromAny(obj,NULL,0,0,0,NULL) calls
+         * PyArray_FromArray(obj,NULL,0) which return original obj
+         */
+        if(PyArray_IsZeroDim(obj) && PyArray_Check(obj)){
+            Py_INCREF(obj);
+            out_op[i] = (PyArrayObject *)obj;
+        }else{
+            if (!PyArray_Check(obj) && !PyArray_IsScalar(obj, Generic)) {
+                /*
+                 * TODO: There should be a comment here explaining what
+                 *       context does.
+                 */
+                context = Py_BuildValue("OOi", ufunc, args, i);
+                if (context == NULL) {
+                    return -1;
+                }
             }
-        }
-        else {
-            context = NULL;
-        }
-        out_op[i] = (PyArrayObject *)PyArray_FromAny(obj,
+            else {
+                context = NULL;
+            }
+            out_op[i] = (PyArrayObject *)PyArray_FromAny(obj,
                                     NULL, 0, 0, 0, context);
-        Py_XDECREF(context);
+            Py_XDECREF(context);
+        }
+
         if (out_op[i] == NULL) {
             return -1;
         }