Tensorflow TF中稀疏张量的einsum

如何解决Tensorflow TF中稀疏张量的einsum

我想将两个张量相乘，一个张量稀疏，另一个张量密集。稀疏的是3D，密集的是2D。我无法将稀疏张量转换为密集张量（即，避免使用tf.sparse.to_dense(...)）。

我的乘法是由以下定律给出的：

C[i,j] = \sum_k A[i,k]*B[i,k,j]

其中C = A * B，A和B是上述的密集和稀疏张量。

在TF中执行的示例如下：

# Example
# inputs
A = tf.constant([[1,2,3],[4,5,6],[7,8,9]],tf.float32)
B = tf.sparse.SparseTensor(values=tf.constant([1,1,-1],tf.float32),indices=[[0,1],[0,2],0],[1,[2,2]],dense_shape=[3,3,3])


# output
C = tf.constant([[3,10,11],[9,-1]],tf.float32)

tf.einsum不支持稀疏张量。

我有一个版本，其中我将3D稀疏张量B切片为2D稀疏矩阵B[0,:,:],B[1,B[2,...的集合，并将密集矩阵A的每一行A[i,:]与每个2D切片稀疏相乘矩阵B[i,:]应用tf.sparse.sparse_dense_matmul(A[i,B[i,:])函数（在切片后具有相应的重塑以将2D张量作为tf.sparse.sparse_dense_matmul的参数）。然后，我堆叠所有矢量结果以组装C矩阵。此过程很慢，并且破坏了B的张量结构。我想通过仅应用Tensorflow函数来执行相同的操作（避免使用循环来对稀疏张量进行切片和破坏，以便稍后通过堆叠重新确定结果）。然后，这应该与Keras作为神经网络的一层一起工作（[A，B]是输入的批处理列表，C = A * B是该层的批处理输出）。打破张量来计算乘法对于在编译图中的训练是疯狂的！

有什么想法吗？稀疏张量是否存在任何类似于tf.sparse.einsum的函数？

如果我将B转换为密集张量，则应用tf.einsum(A,B,'ik,ikj->ij')会非常简单。但是，我不能失去B的稀疏性。

谢谢。问候