PyTorch equivalent to Tensorflow's tf.keras.dot() in terms of behaviour?
In tensorflow, if you have 2 tensors of shape NxTxD and NxDxT respectively (N=batch_size, T=SequenceLength, D=NumberOfFeatures), you can dot them and get an output of NxTxT, as demonstrated below:
import tensorflow as tf
import numpy as np

x1 = np.arange(2 * 4 * 3).reshape(2, 4, 3)
x2 = np.flip(np.arange(2 * 4 * 3).reshape(2, 3, 4), 1).copy()
print(x1.shape, x2.shape)
dotted = tf.keras.layers.Dot(axes=(2, 1))([x1, x2])
print(dotted.shape)
dotted

(2, 4, 3) (2, 3, 4)
(2, 4, 4)
&lt;tf.Tensor: shape=(2, 4, 4), dtype=int32, numpy=
array([[[   4,    7,   10,   13],
        [  40,   52,   64,   76],
        [  76,   97,  118,  139],
        [ 112,  142,  172,  202]],

       [[ 616,  655,  694,  733],
        [ 760,  808,  856,  904],
        [ 904,  961, 1018, 1075],
        [1048, 1114, 1180, 1246]]])&gt;

If you try to do the same in PyTorch, the result is different:
import torch
import numpy as np

x1 = torch.from_numpy(np.arange(2 * 4 * 3).reshape(2, 4, 3))
x2 = torch.from_numpy(np.flip(np.arange(2 * 4 * 3).reshape(2, 3, 4), 1).copy())
dotted = torch.tensordot(x1, x2, dims=([2], [1]))
print(x1.shape, x2.shape)
print(dotted.shape)
dotted

torch.Size([2, 4, 3]) torch.Size([2, 3, 4])
torch.Size([2, 4, 2, 4])
tensor([[[[   4,    7,   10,   13],
          [  40,   43,   46,   49]],

         [[  40,   52,   64,   76],
          [ 184,  196,  208,  220]],

         [[  76,   97,  118,  139],
          [ 328,  349,  370,  391]],

         [[ 112,  142,  172,  202],
          [ 472,  502,  532,  562]]],


        [[[ 148,  187,  226,  265],
          [ 616,  655,  694,  733]],

         [[ 184,  232,  280,  328],
          [ 760,  808,  856,  904]],

         [[ 220,  277,  334,  391],
          [ 904,  961, 1018, 1075]],

         [[ 256,  322,  388,  454],
          [1048, 1114, 1180, 1246]]]], dtype=torch.int32)

Now, Tensorflow's results exist inside the results that pytorch produces (it's a subset of it).
In fact, tensorflow's results is basically some kind of &quot;diagonal&quot; in higher dimensions.
PyTorch's output is NxTxNxT, so to get exactly the same results as Tensorflow you can do:
torch.stack([dotted[i, :, i, :] for i in range(len(dotted))])

tensor([[[   4,    7,   10,   13],
         [  40,   52,   64,   76],
         [  76,   97,  118,  139],
         [ 112,  142,  172,  202]],

        [[ 616,  655,  694,  733],
         [ 760,  808,  856,  904],
         [ 904,  961, 1018, 1075],
         [1048, 1114, 1180, 1246]]], dtype=torch.int32)

but this doesn't negate the fact that you're both:

Allocating memory for a tensor of NxTxNxT instead of NxTxT
The computational complexity/time increases dramatically

Is there a way to get the same 3 dimensional results that tensorflow gives from pytorch, without it computing the 4 dimensional tensor?

I hope you are looking for batch matrix multiplication (bmm) which multiplies two batches of matrices - the two tensors have to be 3D. https://pytorch.org/docs/stable/generated/torch.bmm.html
