是否可以在 Tensorflow 中对 3 阶矩阵执行稀疏-密集矩阵乘法？

Question

I am trying to perform sparse matrix - dense matrix multiplication in TensorFlow, where both matrices have a leading batch dimension (ie, rank 3). I am aware that TensorFlow provides the tf.sparse.sparse_dense_matmul function for rank 2 matrices, but I am looking for a method to handle rank 3 matrices. Is there a built-in function or method in TensorFlow that can handle this case efficiently, without the need for expensive reshaping or slicing operations? Performance is critical in my application.

To illustrate my question, consider the following example code:

import tensorflow as tf

# Define sparse and dense matrices with leading batch dimension
sparse_tensor = tf.SparseTensor(indices=[[0, 1, 1], [0, 0, 1], [1, 1, 1], [1, 2, 1], [2, 1, 1]],
                                values=[1, 1, 1, 1, 1],
                                dense_shape=[3, 3, 2])
dense_matrix = tf.constant([[[0.1, 0.2, 0.3, 0.4], [0.5, 0.6, 0.7, 0.8]],
                        [[0.9, 0.10, 0.11, 0.12], [0.13, 0.14, 0.15, 0.16]],
                        [[0.17, 0.18, 0.19, 0.20], [0.21, 0.22, 0.23, 0.24]]], dtype=tf.float32)
        
# Perform sparse x dense matrix multiplication
result = tf.???(sparse_tensor, dense_matrix)  # Result should have shape [3, 3, 4]

Answer 1

In TF, Sparse and Dense multiplication broadcasts dense to sparse only. Otherwise batch_sparse_dense_matmul can be done simply by,

tf.sparse.reduce_sum(tf.sparse.expand_dims(sparse_tensor,-1)*tf.expand_dims(dense_matrix,1), 2)
#[3,3,2,1] * [3,1,2,4] and reduce sum along dim=2
# the above throws error 
# because the last dim of sparse tensor [1] cannot be broadcasted to [4]

To fix the above issue, we need to tile the last dimension of the sparse tensor to make it 4.

k = 4
tf.sparse.concat(-1,[tf.sparse.expand_dims(sparse_tensor, -1)]*k)
##[3,3,2,4]

Putting together,

tf.sparse.reduce_sum(tf.sparse.concat(-1,[tf.sparse.expand_dims(sparse_tensor, -1)]*k)*tf.expand_dims(dense_matrix,1), 2)
#timeit:1.99ms

Another way would be to use tf.map_fn ,

tf.map_fn(
        lambda x: tf.sparse.sparse_dense_matmul(x[0], x[1]),
        elems=(tf.sparse.reorder(tf.cast(sparse_tensor, tf.float32)),dense_matrix ), fn_output_signature=tf.float32
    )
#timeit:4.42ms