torch.tensordot

torch.tensordot(a, b, dims=2, out=None) [исходный код]

Возвращает результат свёртки a и b по нескольким измерениям.

tensordot реализует обобщённое матричное умножение.

Параметры:

• a (Tensor) – Левый тензор для свёртки
• b (Tensor) – Правый тензор для свёртки
• dims (int or Tuple[List[int], List[int]] or List[List[int]] containing two lists or Tensor) – количество измерений для свёртки или явные списки измерений соответственно для a и b

При вызове с неотрицательным целочисленным аргументом dims = $d$, если количество измерений a и b равно соответственно $m$ и $n$, tensordot() вычисляет тензор $r$ формы a.shape[:-dims] + b.shape[dims:], заданный следующим образом:

$$r_{i_1,...,i_{m-d}, j_1,...,j_{n-d}} = \sum_{k_1,...,k_d} a_{i_1,...,i_{m-d},k_1,...,k_d} \times b_{k_1,...,k_d, j_1,...,j_{n-d}}.$$

При вызове с dims в виде списка указанные измерения будут свёрнуты вместо последних $d$ измерений a и первых $d$ измерений $b$. Размеры по этим измерениям должны совпадать, однако tensordot() обрабатывает измерения с широковещательной передачей.

Примеры:

>>> a = torch.arange(60.).reshape(3, 4, 5)
>>> b = torch.arange(24.).reshape(4, 3, 2)
>>> torch.tensordot(a, b, dims=([1, 0], [0, 1]))
tensor([[4400., 4730.],
        [4532., 4874.],
        [4664., 5018.],
        [4796., 5162.],
        [4928., 5306.]])

>>> a = torch.randn(3, 4, 5, device='cuda')
>>> b = torch.randn(4, 5, 6, device='cuda')
>>> c = torch.tensordot(a, b, dims=2).cpu()
tensor([[ 8.3504, -2.5436,  6.2922,  2.7556, -1.0732,  3.2741],
        [ 3.3161,  0.0704,  5.0187, -0.4079, -4.3126,  4.8744],
        [ 0.8223,  3.9445,  3.2168, -0.2400,  3.4117,  1.7780]])

>>> a = torch.randn(3, 5, 4, 6)
>>> b = torch.randn(6, 4, 5, 3)
>>> torch.tensordot(a, b, dims=([2, 1, 3], [1, 2, 0]))
tensor([[  7.7193,  -2.4867, -10.3204],
        [  1.5513, -14.4737,  -6.5113],
        [ -0.2850,   4.2573,  -3.5997]])