Is sparse tensor multiplication implemented in TensorFlow?

chris picture chris · Dec 1, 2015 · Viewed 13.5k times · Source

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

mrry picture mrry · Dec 2, 2015

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).