Element-wise multiply a dense vector with each row of a sparse matrix in Tensorflow-CodePudding

Suppose I have a sparse matrix

A = tf.sparse.SparseTensor(indices=[[0,0],[1,1],[1,2]], values=[1,1,1], 
                dense_shape=[2,3])

and a dense vector

B = tf.constant([4,3,5])

The shape of matrix A and vector B are (2,3), (1,3) respectively. I would like to element-wise multiply B with each row of A. The expected result is another sparse matrix, say

C = tf.sparse.SparseTensor(indices=[[0,0],[1,1],[1,2]], values=[4,3,5], 
                dense_shape=[2,3])

I know it would be relatively easy if A is a dense matrix, but the dense size of A is extremely large and most of the elements in A are zero.

CodePudding user response：

Just multiplying with an asterisk * works.

tf.reduce_all(tf.sparse.to_dense(A * B) == tf.sparse.to_dense(C))

<tf.Tensor: shape=(), dtype=bool, numpy=True>

Btw, B has shape (3,), not (1, 3)

This is the result of this operation:

<tf.Tensor: shape=(2, 3), dtype=int32, numpy=
array([[4, 0, 0],
       [0, 3, 5]])>

You could also have done this manually, but keep an eye on the broadcasting rules:

tf.sparse.SparseTensor(indices=A.indices, values=A.values * B, 
                       dense_shape=A.dense_shape)