convert CSR format to dense/COO format in tensorflow

Question

tf.sparse_to_dense() fucntion in tensorflow only support ((data, (row_ind, col_ind)), [shape=(M, N)]) format. How can I convert standard CSR tensor (((data, indices, indptr), [shape=(M, N)])) to dense representation in tensorflow?

For example given, data, indices and indptr the function will return dense tensor.

e.g., inputs:

indices = [1 3 3 0 1 2 2 3]
indptr = [0 2 3 6 8]
data = [2 4 1 3 2 1 1 5]

expected output:

[[0, 2, 0, 4],
 [0, 0, 0, 1],
 [3, 2, 1, 0], 
 [0, 0, 1, 5]]

According to Scipy documentation, we can convert it back by the following:

the column indices for row i are stored in indices[indptr[i]:indptr[i+1]] and their corresponding values are stored in data[indptr[i]:indptr[i+1]]. If the shape parameter is not supplied, the matrix dimensions are inferred from the index arrays.

user_7370153 · Answer

It is relatively easily to convert from the CSR format to the COO by expanding the indptr argument to get the row indices. Here is an example using a subtraction, tf.repeat and tf.range. The shape of the final sparse tensor is inferred from the max indices in the rows/columns respectively (but can also be provided explicitly).

def csr_to_sparse(data, indices, indptr, dense_shape=None):
    rep = tf.math.subtract(indptr[1:], indptr[:-1])
    row_indices = tf.repeat(tf.range(tf.size(rep)), rep)
    sparse_indices = tf.cast(tf.stack((row_indices, indices), axis=-1), tf.int64)
    if dense_shape is None:
        max_row = tf.math.reduce_max(row_indices)
        max_col = tf.math.reduce_max(indices)
        dense_shape = (max_row + 1, max_col + 1)
    return tf.SparseTensor(indices=sparse_indices, values=data, dense_shape=dense_shape)

With your example:

>>> indices = [1, 3, 3, 0, 1, 2, 2, 3]
>>> indptr = [0, 2, 3, 6, 8,]
>>> data = [2, 4, 1, 3, 2, 1, 1, 5]
>>> tf.sparse.to_dense(csr_to_sparse(data, indices, indptr))

convert CSR format to dense/COO format in tensorflow

1 Answers