Repository: cupy/cupy
Pull request: Make scipy.special ufuncs work with CuPy inputs
This PR makes it possible to pass CuPy arrays to `scipy.special` ufuncs. The actual implementation is dispatched to `cupyx.scipy.special` using `__array_ufunc__`.

Fixes https://github.com/cupy/cupy/issues/6158.

Problem statement:
Support `scipy.special` with `__array_ufunc__`
### Description

`cupy.ndarray` implements `__array_ufunc__` interface, this allows NumPy ufuncs to accept CuPy arrays and it all just works. The majority of `scipy.special` functions are implemented as ufuncs and the same is true for `cupyx.scipy.special`. It would be great to support interoperability with SciPy's special functions module.

For example, a variant of the following code should work:
```python
import cupy
from scipy.special import rel_entr, kl_div

p = cupy.array([1/2, 1/2])
q = cupy.array([1/10, 9/10])

print(rel_entr(p, q), sum(rel_entr(p, q)))
print(kl_div(p, q), sum(kl_div(p, q)))
assert type(rel_entr(p, q)) == cupy.ndarray
assert type(kl_div(p, q)) == cupy.ndarray
```

Running above code now would give `TypeError: operand type(s) all returned NotImplemented from __array_ufunc__(<ufunc 'rel_entr'>, '__call__', array([0.5, 0.5]), array([0.1, 0.9])): 'ndarray', 'ndarray'`.

Making the above code work correctly is just a two lines of code change:
```python
diff --git a/cupy/core/core.pyx b/cupy/core/core.pyx
index 3d8b27daa..370a143d4 100644
--- a/cupy/core/core.pyx
+++ b/cupy/core/core.pyx
@@ -1317,6 +1317,7 @@ cdef class ndarray:
         numpy array.
         """
         import cupy  # top-level ufuncs
+        import cupyx.scipy.special  # special ufuncs
         if 'out' in kwargs:
             # need to unfold tuple argument in kwargs
             out = kwargs['out']
@@ -1331,7 +1332,7 @@ cdef class ndarray:
                 return NotImplemented
             name = ufunc.__name__
             try:
-                cp_ufunc = getattr(cupy, name)
+                cp_ufunc = getattr(cupy, name, None) or getattr(cupyx.scipy.special, name)
             except AttributeError:
                 return NotImplemented
             if name in [
```

Applying the patch and recompiling CuPy makes the example from the first post work:
```python
In [1]: import cupy
   ...: from scipy.special import rel_entr, kl_div
   ...:
   ...: p = cupy.array([1/2, 1/2])
   ...: q = cupy.array([1/10, 9/10])
   ...:
   ...: print(rel_entr(p, q), sum(rel_entr(p, q)))
   ...: print(kl_div(p, q), sum(kl_div(p, q)))
   ...: assert type(rel_entr(p, q)) == cupy.ndarray
   ...: assert type(kl_div(p, q)) == cupy.ndarray
[ 0.80471896 -0.29389333] 0.5108256237659907
[0.40471896 0.10610667] 0.5108256237659906

In [2]: type(rel_entr(p, q))
Out[2]: cupy.core.core.ndarray
```

Is this something that would be acceptable to have in CuPy?

### Additional Information

_No response_

Hints from issue discussion:
@IvanYashchuk Just for my own curiosity: Do you know how `scipy.special` supports `__array_ufunc__`? Did they use the ufunc C APIs?
Yes, they use ufunc's C API ([`PyUFunc_FromFuncAndData`](https://numpy.org/devdocs/reference/c-api/ufunc.html?highlight=pyufunc_fromfuncanddata#c.PyUFunc_FromFuncAndData)) and they generate Cython code that binds Cython/C/C++/Fortran implementation to Python function.
See here https://github.com/scipy/scipy/blob/b6481610941b62d23bdd9a104840b58da35cac0b/scipy/special/_generate_pyx.py#L773-L777
and the docstring at the top of the file.

Commit message: Fix too long line

Patch to review:
diff --git a/cupy/_core/core.pyx b/cupy/_core/core.pyx
index 9978d796488..35385b48d10 100644
--- a/cupy/_core/core.pyx
+++ b/cupy/_core/core.pyx
@@ -1545,6 +1545,7 @@ cdef class ndarray:
         numpy array.
         """
         import cupy  # top-level ufuncs
+        import cupyx.scipy.special  # special ufuncs
         inout = inputs
         if 'out' in kwargs:
             # need to unfold tuple argument in kwargs
@@ -1559,7 +1560,9 @@ cdef class ndarray:
         if method == '__call__':
             name = ufunc.__name__
             try:
-                cp_ufunc = getattr(cupy, name)
+                cp_ufunc = getattr(cupy, name, None) or getattr(
+                    cupyx.scipy.special, name
+                )
             except AttributeError:
                 return NotImplemented
             for x in inout:
diff --git a/tests/cupyx_tests/scipy_tests/special_tests/test_ufunc_dispatch.py b/tests/cupyx_tests/scipy_tests/special_tests/test_ufunc_dispatch.py
new file mode 100644
index 00000000000..01c76fd18dc
--- /dev/null
+++ b/tests/cupyx_tests/scipy_tests/special_tests/test_ufunc_dispatch.py
@@ -0,0 +1,42 @@
+import numpy
+import cupy
+import cupyx.scipy.special
+
+from cupy import testing
+import pytest
+
+try:
+    import scipy.special
+
+    scipy_ufuncs = {
+        f
+        for f in scipy.special.__all__
+        if isinstance(getattr(scipy.special, f), numpy.ufunc)
+    }
+    cupyx_scipy_ufuncs = {
+        f
+        for f in dir(cupyx.scipy.special)
+        if isinstance(getattr(cupyx.scipy.special, f), cupy.ufunc)
+    }
+except ImportError:
+    scipy_ufuncs = set()
+    cupyx_scipy_ufuncs = set()
+
+
+@testing.gpu
+@testing.with_requires("scipy")
+@pytest.mark.parametrize("ufunc", sorted(cupyx_scipy_ufuncs & scipy_ufuncs))
+class TestUfunc:
+    @testing.numpy_cupy_allclose(atol=1e-5)
+    def test_dispatch(self, xp, ufunc):
+        ufunc = getattr(scipy.special, ufunc)
+        # some ufunc (like sph_harm) do not work with float inputs
+        # therefore we retrieve the types from the ufunc itself
+        types = ufunc.types[0]
+        args = [
+            cupy.testing.shaped_random((5,), xp, dtype=types[i])
+            for i in range(ufunc.nargs - 1)
+        ]
+        res = ufunc(*args)
+        assert type(res) == xp.ndarray
+        return res


Code review comments:
diff_hunk: @@ -0,0 +1,42 @@
+import numpy
+import cupy
+import cupyx.scipy.special
+
+from cupy import testing
+import pytest
+
+try:
+    import scipy.special
+
+    scipy_ufuncs = {
+        f
+        for f in scipy.special.__all__
+        if isinstance(getattr(scipy.special, f), numpy.ufunc)
+    }
+    cupyx_scipy_ufuncs = {
+        f
+        for f in dir(cupyx.scipy.special)
+        if isinstance(getattr(cupyx.scipy.special, f), cupy.ufunc)
+    }
+except ImportError:
+    scipy_ufuncs = set()
+    cupyx_scipy_ufuncs = set()
+
+
+@testing.gpu
+@testing.with_requires("scipy")
+@pytest.mark.parametrize("ufunc", sorted(cupyx_scipy_ufuncs & scipy_ufuncs))
+class TestUfunc:
+    @testing.numpy_cupy_allclose(atol=1e-5)
+    def test_dispatch(self, xp, ufunc):
+        ufunc = getattr(scipy.special, ufunc)
+        # some ufunc (like sph_harm) do not work with float inputs
+        # therefore we retrieve the types from the ufunc itself
+        types = ufunc.types[0]
+        args = [
+            cupy.testing.shaped_random((5,), xp, dtype=types[i])
+            for i in range(ufunc.nargs - 1)
+        ]
+        res = ufunc(*args)
+        assert type(res) == xp.ndarray
+        return res
path: tests/cupyx_tests/scipy_tests/special_tests/test_ufunc_dispatch.py
text: @user1:
```suggestion
        return ufunc(*args)
```
`numpy_cupy_allclose` checks return types.

diff_hunk: @@ -0,0 +1,42 @@
+import numpy
+import cupy
+import cupyx.scipy.special
+
+from cupy import testing
+import pytest
+
+try:
+    import scipy.special
+
+    scipy_ufuncs = {
+        f
+        for f in scipy.special.__all__
+        if isinstance(getattr(scipy.special, f), numpy.ufunc)
+    }
+    cupyx_scipy_ufuncs = {
+        f
+        for f in dir(cupyx.scipy.special)
+        if isinstance(getattr(cupyx.scipy.special, f), cupy.ufunc)
+    }
+except ImportError:
+    scipy_ufuncs = set()
+    cupyx_scipy_ufuncs = set()
+
+
+@testing.gpu
+@testing.with_requires("scipy")
+@pytest.mark.parametrize("ufunc", sorted(cupyx_scipy_ufuncs & scipy_ufuncs))
+class TestUfunc:
+    @testing.numpy_cupy_allclose(atol=1e-5)
+    def test_dispatch(self, xp, ufunc):
+        ufunc = getattr(scipy.special, ufunc)
+        # some ufunc (like sph_harm) do not work with float inputs
+        # therefore we retrieve the types from the ufunc itself
+        types = ufunc.types[0]
+        args = [
+            cupy.testing.shaped_random((5,), xp, dtype=types[i])
+            for i in range(ufunc.nargs - 1)
path: tests/cupyx_tests/scipy_tests/special_tests/test_ufunc_dispatch.py
text: @user1:
```suggestion
            for i in range(ufunc.nin)
```
for readability and to prepare ufuncs like `np.divmod`.