cupy.unique()函数中是否有CuPy版本支持(axis)选项？有什么解决方法吗？

Question

我正在寻找一个支持轴选项的numpy.unique()的图形处理器CuPy对应物。

我有一个Cupy 2D数组，我需要删除它的重复行。不幸的是，cupy.unique()函数将数组展平并返回具有唯一值的一维数组。我正在寻找像numpy.unique(arr，axis=0)这样的函数来解决这个问题，但是CuPy还不支持(axis)选项

x = cp.array([[1,2,3,4], [4,5,6,7], [1,2,3,4], [10,11,12,13]])
y = cp.unique(x)
y_r = np.unique(cp.asnumpy(x), axis=0)

print('The 2D array:\n', x)
print('Required:\n', y_r, 'But using CuPy')
print('The flattened unique array:\n', y)

print('Error producing line:', cp.unique(x, axis=0))

I expect a 2D array with unique rows but I get a 1D array with unique numbers instead. Any ideas about how to implement this with CuPy or numba?

Answer 1

从CuPy版本8.0.0b2开始，函数cupy.lexsort已正确实现。此函数可用作使用axis参数的cupy.unique的变通方法(尽管可能不是最有效的)。

假设数组是2D的，并且您希望找到沿轴0的唯一元素(否则根据需要转置/交换)：

    ###################################
    # replacement for numpy.unique with option axis=0
    ###################################

    def cupy_unique_axis0(array):
        if len(array.shape) != 2:
            raise ValueError("Input array must be 2D.")
        sortarr     = array[cupy.lexsort(array.T[::-1])]
        mask        = cupy.empty(array.shape[0], dtype=cupy.bool_)
        mask[0]     = True
        mask[1:]    = cupy.any(sortarr[1:] != sortarr[:-1], axis=1)
        return sortarr[mask]

如果您还想实现return_stuff参数，请选中the original cupy.unique source code (这是它的基础)。我自己也不需要这些。