CLN: optimize ndarray.__array_function__

1 files changed, 6 insertions, 3 deletions

diff --git a/numpy/core/_methods.py b/numpy/core/_methods.py
index f49f64feb..8974f0ce1 100644
--- a/numpy/core/_methods.py
+++ b/numpy/core/_methods.py
@@ -155,13 +155,16 @@ def _ptp(a, axis=None, out=None, keepdims=False):
         out
     )
 
+_NDARRAY_ARRAY_FUNCTION = mu.ndarray.__array_function__
+
 def _array_function(self, func, types, args, kwargs):
     # TODO: rewrite this in C
     # Cannot handle items that have __array_function__ other than our own.
     for t in types:
-        if (hasattr(t, '__array_function__') and
-                t.__array_function__ is not mu.ndarray.__array_function__):
-            return NotImplemented
+        if t is not mu.ndarray:
+            method = getattr(t, '__array_function__', _NDARRAY_ARRAY_FUNCTION)
+            if method is not _NDARRAY_ARRAY_FUNCTION:
+                return NotImplemented
 
     # Arguments contain no overrides, so we can safely call the
     # overloaded function again.