Scaling columns of a Sparse Tensor by a vector in tensorflow

Question

Assume A is a sparse 2D binary tensor (NxN) and B is a vector (1xN). I want to scale columns of A by B: A_{i, j} <- A_{i, j} x B_{j} I am looking for an efficient way for doing this. Currently I tried two methods but they are too slow:

1. Convert B to the diagonal matrix BD and use tf.sparse_tensor_dense_matmul(A, BD) . This is going to return a dense matrix and as a result, the method is not efficient.
2. Using map_fn for sparse tensors tf.SparseTensor(A.indices, tf.map_fn(func, (A.values, A.indices[:, 1]), dtype=tf.float32), A.dense_shape) . This returns a sparse tensor but it seems one of the map_fn or SparseTensor functions are the bottleneck.

Thank you!

Answer 1

The faster solution is probably to call gather on your tensor B to be able to multiply directly the values of A . Another option is to use tf.sparse.map_values , creating your vector of updates with tf.gather .

On a 99% empty sparse matrix of size 1000x1000 :

>>> %timeit tf.sparse.SparseTensor(A.indices, A.values*tf.gather(B,A.indices[:,1]), A.dense_shape)
855 µs ± 86 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

The map_values approach:

>>> %timeit tf.sparse.map_values(tf.multiply, A, tf.gather(B,A.indices[:,1]))
978 µs ± 73 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Compared to the sparse_dense_matmul approach:

>>> %timeit tf.sparse.from_dense(tf.sparse.sparse_dense_matmul(A,tf.linalg.diag(B)))
26.7 ms ± 2.24 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)