Remove temporary workaround from dpnp.put_along_axis (#1839)

* Get rid of call_origin in dpnp.put

* Removed temporary w/a in dpnp.put_along_axis

Author: antonwolfy

.github/workflows/conda-package.yml

@@ -49,6 +49,7 @@ env:
       test_usm_type.py
       third_party/cupy/core_tests
       third_party/cupy/indexing_tests/test_indexing.py
+      third_party/cupy/lib_tests
       third_party/cupy/linalg_tests
       third_party/cupy/logic_tests
       third_party/cupy/manipulation_tests

dpnp/dpnp_iface_indexing.py

@@ -831,8 +831,7 @@ def put(a, ind, v, /, *, axis=None, mode="wrap"):
         in_a[:] = a.reshape(in_a.shape, copy=False)
 
 
-# pylint: disable=redefined-outer-name
-def put_along_axis(a, indices, values, axis):
+def put_along_axis(a, ind, values, axis):
     """
     Put values into the destination array by matching 1d index and data slices.
 
@@ -842,13 +841,13 @@ def put_along_axis(a, indices, values, axis):
     ----------
     a : {dpnp.ndarray, usm_ndarray}, (Ni..., M, Nk...)
         Destination array.
-    indices : {dpnp.ndarray, usm_ndarray}, (Ni..., J, Nk...)
+    ind : {dpnp.ndarray, usm_ndarray}, (Ni..., J, Nk...)
         Indices to change along each 1d slice of `a`. This must match the
         dimension of input array, but dimensions in ``Ni`` and ``Nj``
         may be 1 to broadcast against `a`.
     values : {scalar, array_like}, (Ni..., J, Nk...)
         Values to insert at those indices. Its shape and dimension are
-        broadcast to match that of `indices`.
+        broadcast to match that of `ind`.
     axis : int
         The axis to take 1d slices along. If axis is ``None``, the destination
         array is treated as if a flattened 1d view had been created of it.
@@ -880,16 +879,12 @@ def put_along_axis(a, indices, values, axis):
 
     """
 
-    dpnp.check_supported_arrays_type(a, indices)
-
-    # TODO: remove when #1382(dpctl) is resolved
-    if dpnp.is_supported_array_type(values) and a.dtype != values.dtype:
-        values = values.astype(a.dtype)
+    dpnp.check_supported_arrays_type(a, ind)
 
     if axis is None:
         a = a.ravel()
 
-    a[_build_along_axis_index(a, indices, axis)] = values
+    a[_build_along_axis_index(a, ind, axis)] = values
 
 
 def putmask(x1, mask, values):

tests/third_party/cupy/lib_tests/test_shape_base.py

@@ -0,0 +1,167 @@
+import unittest
+
+import numpy
+import pytest
+
+import dpnp as cupy
+from tests.third_party.cupy import testing
+
+
+@testing.parameterize(*(testing.product({"axis": [0, 1, -1]})))
+@pytest.mark.skip("'apply_along_axis' is not implemented yet")
+class TestApplyAlongAxis(unittest.TestCase):
+    @testing.numpy_cupy_array_equal()
+    def test_simple(self, xp):
+        a = xp.ones((20, 10), "d")
+        return xp.apply_along_axis(len, self.axis, a)
+
+    @testing.for_all_dtypes(no_bool=True)
+    @testing.numpy_cupy_array_equal()
+    def test_3d(self, xp, dtype):
+        a = xp.arange(27, dtype=dtype).reshape((3, 3, 3))
+        return xp.apply_along_axis(xp.sum, self.axis, a)
+
+    @testing.numpy_cupy_array_equal()
+    def test_0d_array(self, xp):
+
+        def sum_to_0d(x):
+            """Sum x, returning a 0d array of the same class"""
+            assert x.ndim == 1
+            return xp.squeeze(xp.sum(x, keepdims=True))
+
+        a = xp.ones((6, 3))
+        return xp.apply_along_axis(sum_to_0d, self.axis, a)
+
+    @testing.numpy_cupy_array_equal()
+    def test_axis_insertion_2d(self, xp):
+
+        def f1to2(x):
+            """produces an asymmetric non-square matrix from x"""
+            assert x.ndim == 1
+            return x[::-1] * x[1:, None]
+
+        # 2d insertion
+        a2d = xp.arange(6 * 3).reshape((6, 3))
+        return xp.apply_along_axis(f1to2, self.axis, a2d)
+
+    @testing.numpy_cupy_array_equal()
+    def test_axis_insertion_3d(self, xp):
+
+        def f1to2(x):
+            """produces an asymmetric non-square matrix from x"""
+            assert x.ndim == 1
+            return x[::-1] * x[1:, None]
+
+        # 3d insertion
+        a3d = xp.arange(6 * 5 * 3).reshape((6, 5, 3))
+        return xp.apply_along_axis(f1to2, self.axis, a3d)
+
+    def test_empty1(self):
+        # can't apply_along_axis when there's no chance to call the function
+        def never_call(x):
+            assert False  # should never be reached
+
+        for xp in [numpy, cupy]:
+            a = xp.empty((0, 0))
+            with pytest.raises(ValueError):
+                xp.apply_along_axis(never_call, self.axis, a)
+
+    def test_empty2(self):
+        # but it's sometimes ok with some non-zero dimensions
+        def empty_to_1(x):
+            assert len(x) == 0
+            return 1
+
+        for xp in [numpy, cupy]:
+            shape = [10, 10]
+            shape[self.axis] = 0
+            shape = tuple(shape)
+            a = xp.empty(shape)
+            if self.axis == 0:
+                other_axis = 1
+            else:
+                other_axis = 0
+            with pytest.raises(ValueError):
+                xp.apply_along_axis(empty_to_1, other_axis, a)
+
+            # okay to call along the shape 0 axis
+            testing.assert_array_equal(
+                xp.apply_along_axis(empty_to_1, self.axis, a), 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.with_requires("numpy>=1.16")
+@pytest.mark.skip("'apply_along_axis' is not implemented yet")
+def test_apply_along_axis_invalid_axis():
+    for xp in [numpy, cupy]:
+        a = xp.ones((8, 4))
+        for axis in [-3, 2]:
+            with pytest.raises(numpy.AxisError):
+                xp.apply_along_axis(xp.sum, axis, a)
+
+
+class TestPutAlongAxis(unittest.TestCase):
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_put_along_axis_empty(self, xp, dtype):
+        a = xp.array([], dtype=dtype).reshape(0, 10)
+        i = xp.array([], dtype=xp.int64).reshape(0, 10)
+        vals = xp.array([]).reshape(0, 10)
+        ret = xp.put_along_axis(a, i, vals, axis=0)
+        assert ret is None
+        return a
+
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_simple(self, xp, dtype):
+        a = testing.shaped_arange((3, 3, 3), xp, dtype)
+        indices_max = xp.argmax(a, axis=0, keepdims=True)
+        ret = xp.put_along_axis(a, indices_max, 0, axis=0)
+        assert ret is None
+        return a
+
+    @testing.for_all_dtypes()
+    def test_indices_values_arr_diff_dims(self, dtype):
+        for xp in [numpy, cupy]:
+            a = testing.shaped_arange((3, 3, 3), xp, dtype)
+            i_max = xp.argmax(a, axis=0, keepdims=False)
+            with pytest.raises(ValueError):
+                xp.put_along_axis(a, i_max, -99, axis=1)
+
+
+@testing.parameterize(
+    *testing.product(
+        {
+            "axis": [0, 1],
+        }
+    )
+)
+class TestPutAlongAxes(unittest.TestCase):
+    def test_replace_max(self):
+        arr = cupy.array([[10, 30, 20], [60, 40, 50]])
+        indices_max = cupy.argmax(arr, axis=self.axis, keepdims=True)
+        # replace the max with a small value
+        cupy.put_along_axis(arr, indices_max, -99, axis=self.axis)
+        # find the new minimum, which should max
+        indices_min = cupy.argmin(arr, axis=self.axis, keepdims=True)
+        testing.assert_array_equal(indices_min, indices_max)
+
+
+class TestPutAlongAxisNone(unittest.TestCase):
+    @testing.for_all_dtypes()
+    @testing.numpy_cupy_array_equal()
+    def test_axis_none(self, xp, dtype):
+        a = testing.shaped_arange((3, 3), xp, dtype)
+        i = xp.array([1, 3])
+        val = xp.array([99, 100])
+        ret = xp.put_along_axis(a, i, val, axis=None)
+        assert ret is None
+        return a