Computing Linear Layer in Tensor/Outer-Product space in PyTorch is Very Slow
I would like to make a PyTorch model that takes the outer product of the input with itself and then does a linear regression on that. As an example, consider the input vector [1,2,3], then I would like to compute w and b to optimize [1*1, 1*2, 1*3, 2*1, 2*2, 2*3, 3*1, 3*2, 3*3] @ w + b.

For a batch input with r rows and c columns, I can do this in PyTorch with

(input.reshape(r,c,1) @ input.reshape(r,1,c)).reshape(r,c**2) @ weigts + b

My problem is that it is extraordinarily slow. Like a factor 1000 times slower and more memory consumptious than Adding a fully connected c*c RELU layer, even though it has the same number of weights.

My question is why this happens?

Is reshape a very expensive operation for PyTorch? Could I reformulate it in a different way that would make things more efficient?

Another equivalent formulation I know is torch.diag(input @ weights @ input.T) + b, but now we are computing way more values than we need (r*r) just to throw them away again.

When you have to reshape a tensor during the training loop of a model it's always best to use view instead of reshape. There doesn't appear to be any performance overhead with a view, but it does require that the tensor data is contiguous.

If your tensors at the beginning aren't contiguous you can recopy the tensor and make it contiguous.
