我尝试通过矩阵M和X之间的运算创建维度为32 x 576 x 2的输出Tensor，其形状如下：

M.shape: (576, 2, 2048)
X.shape: (32, 2048)

我所定义的运算是一个元素级余弦相似性，由下式得出：

其表示特征向量𝑥和向量M_j，k之间的余弦相似度。
这就是我在代码中实现它的方式（不正确），其中BATCH_SIZE=32，C=576，V=2：

@tf.function
    def call(self, X):
        M = self.kernel

        norm_M = tf.norm(M, ord=2, axis=2)
        norm_X = tf.norm(X, ord=2, axis=1)

        l_r = (some scalar value, separate to this question)

        # Compute cosine similarity between X and M
        # as a matrix with dimensionality:
        # BATCH_SIZE x C x V
        feature_batch_size = tf.shape(X)[0]
        c = tf.shape(M)[0]
        v = tf.shape(M)[1]
        output_matrix = tf.zeros([feature_batch_size, c, v])
        output_matrix = tf.Variable(output_matrix, trainable=False)
        for row in tf.range(feature_batch_size):
          for column in tf.range(c):
            for channel in tf.range(v):
              a = tf.tensordot(M[column][channel], X[row], 1)
              b = norm_M[column][channel] * norm_X[row]
              output_matrix[row][column][channel] = a / b

        return [output_matrix, l_r]

这在output_matrix[row][column][channel] = a / b行上失败，因为它对tf.Variable的单个row：column：channel的赋值不满意。

是否有更好的方法在这两个矩阵上执行此操作以创建所需的输出矩阵，从而可以在不使用这三个嵌套for循环的情况下完成此操作，并保持与tf.Function图形功能的兼容性？
如果没有，我可以做些什么来为tf.Variable上的各个元素赋值，因为我在这里尝试做的不成功？

Extra information:
norm_M.shape: (576, 2)
norm_X.shape: (32,)