Is sparse tensor multiplication implemented in TensorFlow?

Question

Multiplication of sparse tensors with themselves or with dense tensors does not seem to work in TensorFlow. The following example

from __future__ import print_function
import tensorflow as tf

x = tf.constant([[1.0,2.0],
                 [3.0,4.0]])
y = tf.SparseTensor(indices=[[0,0],[1,1]], values=[1.0,1.0], shape=[2,2])
z = tf.matmul(x,y)

sess = tf.Session()
sess.run(tf.initialize_all_variables())
print(sess.run([x, y, z]))

fails with the error message

TypeError: Input 'b' of 'MatMul' Op has type string that does not match type 
float32 of argument 'a'

Both tensors have values of type float32 as seen by evaluating them without the multiplication op. Multiplication of y with itself returns a similar error message. Multipication of x with itself works fine.

Answer 1

General-purpose multiplication for tf.SparseTensor is not currently implemented in TensorFlow. However, there are three partial solutions, and the right one to choose will depend on the characteristics of your data:

• If you have a tf.SparseTensor and a tf.Tensor, you can use tf.sparse_tensor_dense_matmul() to multiply them. This is more efficient than the next approach if one of the tensors is too large to fit in memory when densified: the documentation has more guidance about how to decide between these two methods. Note that it accepts a tf.SparseTensor as the first argument, so to solve your exact problem you will need to use the adjoint_a and adjoint_b arguments, and transpose the result.

• If you have two sparse tensors and need to multiply them, the simplest (if not the most performant) way is to convert them to dense and use tf.matmul:

  a = tf.SparseTensor(...)
  b = tf.SparseTensor(...)
  
  c = tf.matmul(tf.sparse_tensor_to_dense(a, 0.0),
                tf.sparse_tensor_to_dense(b, 0.0),
                a_is_sparse=True, b_is_sparse=True)
  

  Note that the optional a_is_sparse and b_is_sparse arguments mean that "a (or b) has a dense representation but a large number of its entries are zero", which triggers the use of a different multiplication algorithm.

• For the special case of sparse vector by (potentially large and sharded) dense matrix multiplication, and the values in the vector are 0 or 1, the tf.nn.embedding_lookup operator may be more appropriate. This tutorial discusses when you might use embeddings and how to invoke the operator in more detail.

• For the special case of sparse matrix by (potentially large and sharded) dense matrix, tf.nn.embedding_lookup_sparse() may be appropriate. This function accepts one or two tf.SparseTensor objects, with sp_ids representing the non-zero values, and the optional sp_weights representing their values (which otherwise default to one).

Answer 2

Recently, tf.sparse_tensor_dense_matmul(...) was added that allows multiplying a sparse matrix by a dense matrix.

https://www.tensorflow.org/versions/r0.9/api_docs/python/sparse_ops.html#sparse_tensor_dense_matmul

https://github.com/tensorflow/tensorflow/issues/1241

Answer 3

It seems that

tf.sparse_matmul(
    a,
    b,
    transpose_a=None,
    transpose_b=None,
    a_is_sparse=None,
    b_is_sparse=None,
    name=None
)

is not for multiplication of two SparseTensors.

a and b are Tensors not SparseTensors. And I have tried that, it is not working with SparseTensors.

Answer 4

In TF2.4.1 you can use the methods in tensorflow.python.ops.linalg.sparse.sparse_csr_matrix_ops to multiply to arbitrary SparseTensor (I think up to 3 dimensions).

Something like the following should be used (in general you turn the sparse tensors into a CSR representation)

import tensorflow as tf
from tensorflow.python.ops.linalg.sparse import sparse_csr_matrix_ops


def tf_multiply(a: tf.SparseTensor, b: tf.SparseTensor):
    a_sm = sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
        a.indices, a.values, a.dense_shape
    )

    b_sm = sparse_csr_matrix_ops.sparse_tensor_to_csr_sparse_matrix(
        b.indices, b.values, b.dense_shape
    )

    c_sm = sparse_csr_matrix_ops.sparse_matrix_sparse_mat_mul(
        a=a_sm, b=b_sm, type=tf.float32
    )

    c = sparse_csr_matrix_ops.csr_sparse_matrix_to_sparse_tensor(
        c_sm, tf.float32
    )

    return tf.SparseTensor(
        c.indices, c.values, dense_shape=c.dense_shape
    )

For a while I was prefering scipy multiplication (via a py_function) because this multiplication in TF (2.3 and 2.4) was not performing as well as scipy. I tried again recently and, either I changed something on my code, or there was some fix in 2.4.1 that makes the TF sparse multiplication faster that using scipy, in both CPU and GPU.

Answer 5

tf.sparse_matmul is for multiplying two dense tensor not sparse type of data structure. That function is just an optimized version of tensor multiplication if the given matrix (or both two matrixes) have many zero value. Again it does not accept sparse tensor data type. It accepts dense tensor data type. It might fasten your calculations if values are mostly zero.

as far as I know there is no implementation of two sparse type tensor muliplication. but just one sparse one dense which is tf.sparse_tensor_dense_matmul(x, y) !

Answer 6

To make the answer more complete:

tf.sparse_matmul(
    a,
    b,
    transpose_a=None,
    transpose_b=None,
    a_is_sparse=None,
    b_is_sparse=None,
    name=None
)

exists as well:

https://www.tensorflow.org/api_docs/python/tf/sparse_matmul