Slice a list with two other lists in tensorflow / pytorch
How can I slice a list with two other lists? In another word, how can I do a vectorized slicing in tensorflow?
indptr = [0 2 2 5 7]
values = [2 4 3 2 1 1 5]

values[indptr[:-1]:indptr[1:]] # --&gt; throws exception

expected output:
[[2, 4],
 [],
 [3, 2, 1],
 [1, 5]]

More specifically, I wanna vectorized the following loop:
import numpy as np

# sparse representation in CSR format
indptr = [0, 2, 2, 5, 7]
indices = [1, 3, 0, 1, 2, 2, 3]
values = [2, 4, 3, 2, 1, 1, 5]

m, n = 4, 4
out = np.zeros((m, n))

for i in range(m):
    out[i][indices[indptr[i]:indptr[i + 1]]] = values[indptr[i]:indptr[i + 1]]

expected output:
# dense representation
[[0, 2, 0, 4],
 [0, 0, 0, 0],
 [3, 2, 1, 0],
 [0, 0, 1, 5]])


In tensorflow, tf.scatter_nd can be used for the purpose.
@tf.function
def csr_to_dense(indptr,indices,values,m,n):
  repeats=indptr[1:]-indptr[:-1]
  ind1=tf.repeat(tf.range(m),repeats)
  indices=tf.stack([ind1,indices],1)
  return tf.scatter_nd(indices,values,(m,n))

indptr=tf.constant([0,2,2,5,7])
indices=tf.constant([1,3,0,1,2,2,3])
values=tf.constant([2,4,3,2,1,1,5],dtype=tf.float32)
m=tf.constant(4)
n=tf.constant(4)
print(csr_to_dense(indptr,indices,values,m,n))

#from scipy docs, &quot;Duplicate entries are summed together&quot;
indptr=tf.constant([0,9,9,9,9,9])
indices=tf.constant([0,0,0,0,0,0,0,1,2])
values=tf.constant([2,4,3,2,1,-1,5,-1,-3],dtype=tf.float32)
m=tf.constant(5)
n=tf.constant(3)
print(csr_to_dense(indptr,indices,values,m,n))

target=csr_matrix(np.random.randint(9,size=(5,4)))
indptr=tf.constant(target.indptr)
indices=tf.constant(target.indices)
values=tf.constant(target.data,dtype=tf.float32)
m=tf.constant(target.shape[0])
n=tf.constant(target.shape[1])
print(np.allclose(csr_to_dense(indptr,indices,values,m,n),target.toarray()))
'''
tf.Tensor(
[[0. 2. 0. 4.]
 [0. 0. 0. 0.]
 [3. 2. 1. 0.]
 [0. 0. 1. 5.]], shape=(4, 4), dtype=float32)
tf.Tensor(
[[16. -1. -3.]
 [ 0.  0.  0.]
 [ 0.  0.  0.]
 [ 0.  0.  0.]
 [ 0.  0.  0.]], shape=(5, 3), dtype=float32)
True
'''

