python - 3D SparseTensor matrix multiplication with 2D Tensor :InvalidArgumentError: Tensor 'a_shape' must have 2 elements [Op:SparseTensorDenseMatMul]

Question

Hi I am tryigng to do a 3D SparseTensor matrix multiplication with 2D Tensor. Here is a toy example:

3D tensor matrix multiplication with 2D Tensor

import tensorflow as tf
import numpy as np

a = np.array([[[1., 0., 2., 0.],
              [3., 0., 0., 4.]]])
b = (np.array([1., 2.])[:,np.newaxis]).T

a_t = tf.constant(a)
b_t = tf.constant(b)
    
tf.matmul(b_t,a_t)

<tf.Tensor: shape=(1, 1, 4), dtype=float64, numpy=array([[[7., 0., 2., 8.]]])>

3D SparseTensor matrix multiplication with 2D Tensor

import tensorflow as tf
import numpy as np

a = np.array([[[1., 0., 2., 0.],
              [3., 0., 0., 4.]]])
b = (np.array([1., 2.])[:,np.newaxis]).T

a_t = tf.constant(a)
b_t = tf.constant(b)

a_s = tf.sparse.from_dense(a_t)

tf.sparse.sparse_dense_matmul(b_t,a_s)

InvalidArgumentError: Tensor 'a_shape' must have 2 elements [Op:SparseTensorDenseMatMul]

Can you please help me to sort out this error?

Answer 1

The documentation says that the rank of the dense matrix must be 2.

Multiply SparseTensor (or dense Matrix) (of rank 2) "A" by dense matrix

Te first argument of the call of this function is:

SparseTensor (or dense Matrix) A, of rank 2.

So you can do what you're trying to do, at least not this way. Seems like you don't really need the three dimensions, as one of them is just one. If you squeeze it out, it will work:

a_t = tf.constant(tf.squeeze(a))
b_t = tf.constant(b)

a_s = tf.sparse.from_dense(a_t)

tf.sparse.sparse_dense_matmul(b_t,a_s)

<tf.Tensor: shape=(1, 4), dtype=float64, numpy=array([[7., 0., 2., 8.]])>