如何在tensorflow中执行具有不同秩和外维的张量的三对角矩阵的乘法

Question

下面的代码(修改后的tensorflow示例)产生错误“所有输入张量必须具有相同的等级。”tf.linalg.LinearOperatorTridiag的多项运算也给出了类似的误差。我需要在Keras层中将输入乘以三对角矩阵，由于该层输入中的额外批量维度，张量的等级会有所不同。有什么已知的实际解决方案吗？

import tensorflow as tf

superdiag = tf.constant([-1, -1, 0], dtype=tf.float64)
maindiag = tf.constant([2, 2, 2], dtype=tf.float64)
subdiag = tf.constant([0, -1, -1], dtype=tf.float64)
diagonals = [superdiag, maindiag, subdiag]
rhs = tf.constant([[[1, 1], [1, 1], [1, 1]]], dtype=tf.float64)
x = tf.linalg.tridiagonal_matmul(diagonals, rhs, diagonals_format='sequence')

Answer 1

你必须扩展第一个维度

superdiag = tf.constant([-1, -1, 0], dtype=tf.float64)
maindiag = tf.constant([2, 2, 2], dtype=tf.float64)
subdiag = tf.constant([0, -1, -1], dtype=tf.float64)
diagonals = [tf.expand_dims(superdiag,0), tf.expand_dims(maindiag,0), tf.expand_dims(subdiag,0)]
rhs = tf.constant([[[1, 1], [1, 1], [1, 1]]], dtype=tf.float64)
x = tf.linalg.tridiagonal_matmul(diagonals, rhs, diagonals_format='sequence')