Add sparse tensor support to cp.ElementwiseKernel

Description

Hello, I would appreciate if one could use the cp.ElementwiseKernel (and all the other custom kernels) also for sparse tensors.

Additional Information

Example:

import cupy
import cupyx

# Coo Matrix
data  = cupy.array([1, 1, 1, 1, 1, 1], 'f')
row  = cupy.array([0, 1, 1, 2, 2, 2], 'i')
col   = cupy.array([0, 0, 1, 0, 1, 2], 'i')

A     = cupyx.scipy.sparse.coo_matrix((data, (row, col)), shape=(3, 3),dtype=cupy.float32 )
B     = cupyx.scipy.sparse.csr_matrix((data, (row, col)), shape=(3, 3) ,dtype=cupy.float32)
C     = cupyx.scipy.sparse.csc_matrix((data, (row, col)), shape=(3, 3) ,dtype=cupy.float32)

squared_diff = cp.ElementwiseKernel(
   'float32 x, float32 y',
   'float32 z',
   'z = (x - y) * (x - y)',
   'squared_diff')

# Elementwise Kernel for Sparse Matrices.
# Same for A,A ; B,B; C,C
squared_diff(C,C)
# (C-C)*(C-C) works work.

[leofang]

I modified my example above as I see what @takagi had in mind. The old example x-y needs great care because the sparse matrices x and y may not share the same indices for their nonzeros, meaning you can’t just naively run the x-y ElementwiseKernel on x and y’s data.

[leofang]

Perhaps you can just call arr.eliminate_zeros() after ElementwiseKernel.