Repository: numpy/numpy
Pull request: BUG: Ensure equality/identity comparison with `__array_function__`
Special-casing is needed to ensure the correct public API is passed to `__array_function__` when `like=` kwarg is present.

Fixes #21033 .

Problem statement:
BUG: `np.full` with `like` argument does not trigger `__array_function__` interface
### Describe the issue:

Calling `np.full(shape, x, like=x)` seems to try to trigger the __array_function__ interface, but is not found in the class that implements it. Am I missing something ?

### Reproduce the code example:

```python
import numpy as np
class Toto():
    
    def __init__(self, value):
        self.value = value
        
    def __array_function__(self, func, types, args, kwargs):
        if func not in HANDLED_FUNCTIONS:
            return NotImplemented
        return HANDLED_FUNCTIONS[func](*args, **kwargs)
    
HANDLED_FUNCTIONS = {}

def wrapped_sum(x, *args, **kwargs):
    return Toto(np.sum(x.value, *args, **kwargs))

def wrapped_full(shape, x, *args, **kwargs):
    return Toto(np.full(shape, x.value, *args, **kwargs))

HANDLED_FUNCTIONS[np.sum] = wrapped_sum
HANDLED_FUNCTIONS[np.full] = wrapped_full

print(np.sum(Toto([1, 2, 3]))) # works as expected
print(np.full((3,3), Toto(28), like=Toto(28))) # raise TypeError no implementation found for 'numpy.full' on types that implement __array_function__: [<class '__main__.Toto'>]
```


### Error message:

```shell
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
/var/folders/5k/bf4syt7x1zjbhc6b28srzzym0000gn/T/ipykernel_57970/4259881689.py in <module>
     28 
     29 print(np.sum(Toto([1, 2, 3])))
---> 30 print(np.full((3,3), Toto(28), like=Toto(28)))

/opt/anaconda3/lib/python3.9/site-packages/numpy/core/numeric.py in full(shape, fill_value, dtype, order, like)
    336     """
    337     if like is not None:
--> 338         return _full_with_like(shape, fill_value, dtype=dtype, order=order, like=like)
    339 
    340     if dtype is None:

<__array_function__ internals> in full(*args, **kwargs)

TypeError: no implementation found for 'numpy.full' on types that implement __array_function__: [<class '__main__.Toto'>]
```


### NumPy/Python version information:

1.21.4 3.9.7 (default, Sep 16 2021, 08:50:36) 
[Clang 10.0.0 ]

Hints from issue discussion:
Weird that this wasn't noticed before.  It looks like it passes the wrong public symbol (not `np.full` itself), @pentschev do you have some time to have a look?  Or maybe someone else wants to dig in.  Should not be a hard fix, although the `__array_function__` machinery can confusing.
NumPy's `__array_function__` implementation doesn't guarantee that the `func` argument is the same as the `np.some_function`, i.e., it may wrap a function and thus result in a different object. In your case you're using `np.full` as the key for `HANDLED_FUNCTIONS`, and the function `func` is not the same object. I suggest using the name string (`"full"`, `"sum"`, etc.) as the key in `HANDLED_FUNCTIONS`. The following modification resolves the issue:

```patch
--- numpy-21033.py.orig 2022-02-14 08:11:29.148969000 -0800
+++ numpy-21033.py      2022-02-14 08:12:06.496256000 -0800
@@ -5,9 +5,9 @@
         self.value = value

     def __array_function__(self, func, types, args, kwargs):
-        if func not in HANDLED_FUNCTIONS:
+        if func.__name__ not in HANDLED_FUNCTIONS:
             return NotImplemented
-        return HANDLED_FUNCTIONS[func](*args, **kwargs)
+        return HANDLED_FUNCTIONS[func.__name__](*args, **kwargs)

 HANDLED_FUNCTIONS = {}

@@ -17,8 +17,8 @@
 def wrapped_full(shape, x, *args, **kwargs):
     return Toto(np.full(shape, x.value, *args, **kwargs))

-HANDLED_FUNCTIONS[np.sum] = wrapped_sum
-HANDLED_FUNCTIONS[np.full] = wrapped_full
+HANDLED_FUNCTIONS["sum"] = wrapped_sum
+HANDLED_FUNCTIONS["full"] = wrapped_full

 print(np.sum(Toto([1, 2, 3]))) # works as expected
 print(np.full((3,3), Toto(28), like=Toto(28))) # raise TypeError no implementation found for 'numpy.full' on types that implement __array_function__: [<class '__main__.Toto'>]
```

I suggest also taking a look at [CuPy's `__array_function__` implementation](https://github.com/cupy/cupy/blob/69f9d1111f643c2a663669ddaeaac82e4f665581/cupy/_core/core.pyx#L1606-L1619) as a reference for implementation.
>  i.e., it may wrap a function and thus result in a different object.

Hmmm, can't we just guarantee this, at least so long there was no monkeypatching or anything else weird going on?  It frankly seems strange to me that the publicy API symbol would mismatch, and I was pretty sure the NEP said it should match.
Thanks for the tip, I guess I'll use string lookup from now on.

Also, not sure if related but calling `np.full` without `like=Toto(28)` doesn't seem to trigger the `__array_function__` interface : 
```python
print(np.full((3,3), Toto(28)))
```
returns : 
```

[[<__main__.Toto object at 0x10d1b89d0>
  <__main__.Toto object at 0x10d1b89d0>
  <__main__.Toto object at 0x10d1b89d0>]
 [<__main__.Toto object at 0x10d1b89d0>
  <__main__.Toto object at 0x10d1b89d0>
  <__main__.Toto object at 0x10d1b89d0>]
 [<__main__.Toto object at 0x10d1b89d0>
  <__main__.Toto object at 0x10d1b89d0>
  <__main__.Toto object at 0x10d1b89d0>]]
```
You have to use `like=`, the value doesn't matter.

I personally do think we should fix this though to match up, unless there is some reason why this is impossible.
@seberg I think we're both partially wrong. The NEP-18 doesn't seem to make any claims in either direction that `func` may or may not match `np.some_function`, although the [`example project in NEP-18`](https://numpy.org/neps/nep-0018-array-function-protocol.html#example-for-a-project-implementing-the-numpy-api) suggests the `func` should indeed be the same as `np.some_function`, which points to you being right. However, let me add that I don't think any projects implement it like that, CuPy and Dask certainly don't, which may be an indication of relying on something other than the name being not necessarily the safest.

With that said, I did a bit of digging, and I didn't find any of the functions supported by NEP-18 (pre-`like=`) that `func` didn't match implementation, although I didn't check all of them to verify. For the NEP-35 (post-`like=`), I found that some do match and others don't. Match:

- `np.array`
- `np.asarray`
- `np.asanyarray`
- `np.ascontiguousarray`
- `np.asfortranarray`
- `np.empty`
- `np.zeros`
- `np.arange`
- `np.frombuffer`
- `np.fromiter`
- `np.fromstring`

Do not match:

- `np.require`
- `np.full`
- `np.ones`
- `np.loadtxt`
- `np.genfromtxt`

Also before NumPy 1.21.0, `np.asarray`, `np.asanyarray`, `np.ascontiguousarray` and `np.asfortranarray` didn't match either, but since 1.21.0 they show up as `built-in`, so the implementation that is shown is exposed from C, which wasn't the case before.

I don't immediately see what's the cause for this other than the C vs Python implementation. Please let me know if something comes to mind.
> I don't immediately see what's the cause for this other than the C vs Python implementation. Please let me know if something comes to mind.

Yeah, that looks like C vs. Python.  Can we fix the Python pattern?
After some careful look at the code, I'm not sure we can fix that without moving the implementation to C or rethinking how `like=` dispatching is done. The reason being that when `like=` is specified, we have a new function that is interpreted as the `public_api`, in the case of `np.full` that is [`_full_with_like`](https://github.com/numpy/numpy/blob/5b9b9395036702a23275592293c6501f423269ba/numpy/core/numeric.py#L348-L350) and that's what is seen as `func` by the downstream `__array_function__` implementation, whereas `np.full` becomes the `implementation` in [`array_function_dispatch`'s decorator](https://github.com/numpy/numpy/blob/5b9b9395036702a23275592293c6501f423269ba/numpy/core/overrides.py#L170-L191). This is intentional in the dispatching mechanism of functions with `like=`, as exemplified in [`full`](https://github.com/numpy/numpy/blob/main/numpy/core/numeric.py#L338), but we of course have this inconsistency which I never even thought of given downstream libraries use function names instead of the actual function as key, but NEP-18 originally implemented it in such a way that [testing for equality is valid](https://github.com/numpy/numpy/blob/5b9b9395036702a23275592293c6501f423269ba/numpy/core/src/multiarray/arrayfunction_override.c#L285-L291).

To be fair, I'm not great at thinking in terms of decorators so I won't discard the possibility that there's a simple way to resolve this but I'm failing to see it. I'll give it some more thought this week, but if you happen think of something obvious please let me know @seberg .
@pentschev it shouldn't be hard to fix.  My first angle would be passing `public_api=None` and renaming the function in that scope.  If it is None, replace `public_api` with the new function in the decorator (the inner function).
I don't think we can pass `public_api=None` as this isn't exposed as an argument, but I'm not sure whether you're suggesting we add that somewhere. Renaming the function is done successfully already, but this isn't sufficient for equality testing. For example, printing `public_api` (translates to `_full_with_like`) and `implementation` (translates to `np.full`) we see:

```
public_api: <function full at 0x7f604524f5e0>, implementation: <function full at 0x7f604524f3a0>
```

Perhaps I'm misunderstanding your suggestion?

Commit message: TST: Add check for `__array_function__` equality with `like=` kwarg

Patch to review:
diff --git a/numpy/core/src/multiarray/arrayfunction_override.c b/numpy/core/src/multiarray/arrayfunction_override.c
index 463a2d4d8724..f74e5664a263 100644
--- a/numpy/core/src/multiarray/arrayfunction_override.c
+++ b/numpy/core/src/multiarray/arrayfunction_override.c
@@ -333,7 +333,8 @@ NPY_NO_EXPORT PyObject *
 array_implement_array_function(
     PyObject *NPY_UNUSED(dummy), PyObject *positional_args)
 {
-    PyObject *implementation, *public_api, *relevant_args, *args, *kwargs;
+    PyObject *res, *implementation, *public_api, *relevant_args, *args, *kwargs;
+    int like_removed = 0;
 
     if (!PyArg_UnpackTuple(
             positional_args, "implement_array_function", 5, 5,
@@ -357,11 +358,24 @@ array_implement_array_function(
             }
             Py_DECREF(tmp_has_override);
             PyDict_DelItem(kwargs, npy_ma_str_like);
+            like_removed = 1;
         }
     }
 
-    PyObject *res = array_implement_array_function_internal(
-        public_api, relevant_args, args, kwargs);
+    /*
+     * If `like=` kwarg was removed, `implementation` points to the NumPy
+     * public API, as `public_api` is in that case the wrapper dispatcher
+     * function. For example, in the `np.full` case, `implementation` is
+     * `np.full`, whereas `public_api` is `_full_with_like`. This is done
+     * to ensure `__array_function__` implementations can do equality/identity
+     * comparisons when `like=` is present.
+     */
+    if (like_removed)
+        res = array_implement_array_function_internal(
+            implementation, relevant_args, args, kwargs);
+    else
+        res = array_implement_array_function_internal(
+            public_api, relevant_args, args, kwargs);
 
     if (res == Py_NotImplemented) {
         return PyObject_Call(implementation, args, kwargs);
diff --git a/numpy/core/tests/test_overrides.py b/numpy/core/tests/test_overrides.py
index 9216a3f5fdfa..36970dbc02ed 100644
--- a/numpy/core/tests/test_overrides.py
+++ b/numpy/core/tests/test_overrides.py
@@ -437,6 +437,7 @@ def __init__(self, function=None):
                 self.function = function
 
             def __array_function__(self, func, types, args, kwargs):
+                assert func is getattr(np, func.__name__)
                 try:
                     my_func = getattr(self, func.__name__)
                 except AttributeError:


Code review comments:
diff_hunk: @@ -357,11 +358,24 @@ array_implement_array_function(
             }
             Py_DECREF(tmp_has_override);
             PyDict_DelItem(kwargs, npy_ma_str_like);
+            like_removed = 1;
         }
     }
 
-    PyObject *res = array_implement_array_function_internal(
-        public_api, relevant_args, args, kwargs);
+    /*
+     * If `like=` kwarg was removed, `implementation` points to the NumPy
+     * public API, as `public_api` is in that case the wrapper dispatcher
+     * function. For example, in the `np.full` case, `implementation` is
+     * `np.full`, whereas `public_api` is `_full_with_like`. This is done
+     * to ensure `__array_function__` implementations can do equality/identity
+     * comparisons when `like=` is present.
+     */
+    if (like_removed)
+        res = array_implement_array_function_internal(
+            implementation, relevant_args, args, kwargs);
path: numpy/core/src/multiarray/arrayfunction_override.c
text: @user1:
Could you add braces, we always put them in NumPy.

This is a surprisingly small change!  Could also use `public_api = implementation` with a comment above, but this is OK with me.  I do wonder a bit if it might not be nicer to handle this differently, but so long we have the special case here to drop `like`, I agree that this is completely fine.

Thanks!

@author:
Ah yes, that's a much cleaner way to write that, done in https://github.com/numpy/numpy/pull/21071/commits/059e2b7f5bf1108df6f1f05ad2055c02c9995e2f.

As to handle it differently, maybe would be better to have a different way of handling it, but I would prefer for now a minimal change (such as this one) than potentially having to refactor some of the implementation functions to account for that as it's more likely that we break something else in that case. If you can think of a cleaner solution that won't require much change that can potentially break something else, please let me know.