Repository: cupy/cupy

Issue:
cupy.apply_along_axis failed with cupy.nonzero
### Description

cp.apply_along_axis(cp.nonzero, 1, cp.array([[1,2],[2,3]])) - failed with error

*** AttributeError: 'tuple' object has no attribute 'shape'

np.apply_along_axis(np.nonzero, 1, np.array([[1,2],[2,3]])) - is OK

UPDATE. Problem in _shape_base.py. 
line 53:
buff = cupy.empty(inarr_view.shape[:-1] + res.shape, res.dtype)

res - is tuple (as a result of cp.nonzero(1d array) of single cupy-array, so line 44 ( if cupy.isscalar(res):) doesnt convert it from tuple to cupy-array

as a temporal solution is possible to use "buffer-like" function
def cupy_nonzero (a):
    return cp.nonzero(a)[0]

### To Reproduce

```py
cp.apply_along_axis(cp.nonzero, 1, cp.array([[1,2],[2,3]]))
```


### Installation

_No response_

### Environment

```
# Paste the output here
```


### Additional Information

_No response_


Discussion:
@Satori1313 If the number of nonzero elements in each axis is different, as in the following example, it will not work even if you use NumPy.

```py
>>> np.apply_along_axis(np.nonzero, 1, np.array([[1,0],[2,3]]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 6, in apply_along_axis
  File "***/numpy/lib/shape_base.py", line 402, in apply_along_axis
    buff[ind] = asanyarray(func1d(inarr_view[ind], *args, **kwargs))
ValueError: could not broadcast input array from shape (1,2) into shape (1,1)
```

Do you mean that you want it to work for inputs that you know have the same number of nonzero elements?
@asi1024 Yes, of course))) 
Only if we have the same number of nonzero elements along axis.
@Satori1313 Thanks for your reporting! I will make a fix soon.

However, please note that `cupy.apply_along_axis` is supported only for interface compatibility with NumPy. What `cupy.apply_along_axis` do is just call the given function multiple times without parallelization.

https://github.com/cupy/cupy/blob/d588608a45e14d41bda6057cdb23daef3e40ae85/cupy/lib/_shape_base.py#L57-L58

How about the following algorithm that does not use `apply_along_axis`, to fully use the power of your GPU?

```py
>>> x = cupy.array([[1, 0, 0, 1, 1], [0, 1, 1, 0, 1]])
>>> cupy.nonzero(x)[1].reshape(x.shape[0], -1)
array([[0, 3, 4],
       [1, 2, 4]])
```
@asi1024 Thank you very much!!! I miss that fact! 
And you really help me! Regarding you I can rewrite my func 2x faster!! Thanks again!  

Resolving patch:
diff --git a/cupy/lib/_shape_base.py b/cupy/lib/_shape_base.py
--- a/cupy/lib/_shape_base.py
+++ b/cupy/lib/_shape_base.py
@@ -42,9 +42,7 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
             'Cannot apply_along_axis when any iteration dimensions are 0'
         )
     res = func1d(inarr_view[ind0], *args, **kwargs)
-    if cupy.isscalar(res):
-        # scalar outputs need to be transfered to a device ndarray
-        res = cupy.asarray(res)
+    res = cupy.asarray(res)
 
     # build a buffer for storing evaluations of func1d.
     # remove the requested axis, and add the new ones on the end.
@@ -55,7 +53,8 @@ def apply_along_axis(func1d, axis, arr, *args, **kwargs):
     # save the first result, then compute and save all remaining results
     buff[ind0] = res
     for ind in inds:
-        buff[ind] = func1d(inarr_view[ind], *args, **kwargs)
+        out = func1d(inarr_view[ind], *args, **kwargs)
+        buff[ind] = cupy.asarray(out)
 
     # restore the inserted axes back to where they belong
     for i in range(res.ndim):


Test patch:
diff --git a/tests/cupy_tests/lib_tests/test_shape_base.py b/tests/cupy_tests/lib_tests/test_shape_base.py
--- a/tests/cupy_tests/lib_tests/test_shape_base.py
+++ b/tests/cupy_tests/lib_tests/test_shape_base.py
@@ -91,6 +91,14 @@ def empty_to_1(x):
                 xp.ones((10,))
             )
 
+    @testing.numpy_cupy_array_equal()
+    def test_tuple_outs(self, xp):
+        def func(x):
+            return x.sum(axis=-1), x.prod(axis=-1), x.max(axis=-1)
+
+        a = testing.shaped_arange((2, 2, 2), xp, cupy.int64)
+        return xp.apply_along_axis(func, 1, a)
+
 
 @testing.gpu
 @testing.with_requires('numpy>=1.16')