TensorFlow 中的张量 n 模式积

Question

张量 和矩阵 之间的张量 n 模式乘积定义为 here（第 2.5 节）和 here，用 表示，大小为 。

Tensorflow 中有实现这个操作的函数吗？如果没有，如何使用当前的 API 来实现？到目前为止，我对有关此问题的答案的搜索没有成功。

理想情况下，我会设想一个以 X 和 U 以及 n 作为参数并返回它们对应的 n 模式乘积的函数。

Answer 1

您可以使用tf.einsum 轻松实现固定数量的维度。如果至少 n 是静态已知的（不是符号张量，无论如何这在急切模式下都不应该成为问题），您可以摆弄字符串以获得通用版本：

import tensorflow as tf

def n_mode_product(x, u, n):
    n = int(n)
    # We need one letter per dimension
    # (maybe you could find a workaround for this limitation)
    if n > 26:
        raise ValueError('n is too large.')
    ind = ''.join(chr(ord('a') + i) for i in range(n))
    exp = f'{ind}K...,JK->{ind}J...'
    return tf.einsum(exp, x, u)

# Test
x = tf.ones((2, 3, 4, 5))
u = tf.ones((6, 4))
n = 2  # n is zero-based here
out = n_mode_product(x, u, n)
print(out.shape)
# (2, 3, 6, 5)

您也可以通过将 X 的第 n 轴移动到末尾，然后使用 U 进行（批量）矩阵乘积来获得相同的结果^T 并最终将新的最后一个维度返回到第 n 个位置，但我认为写起来不会更简单也不会更快。