I’m struggling to write a function ConvToToeplitzTr() which takes a conv2d’s weight matrix W as the input, constructs a new matrix using its entries (a Toeplitz matrix representation of the given conv2d, transposed) by reference and returns this new matrix M. Both W and M are then used to calculate the output of the layer.
My goal is for autograd to actually treat the variables in the original conv2d’s weight matrix W, as well as the variables in the output matrix of ConvToToeplitzTr() M as the same and correspondingly accumulate the gradients for each of these variables from both W and M.
The following minimal example already fails to return a tensor by reference:
import torch

def foo(x):
    return torch.stack([x,x])

y = torch.tensor(4)
z = foo(y)
print(z) # [4, 4]
y += 1
print(z) # [4, 4]. Want: [5,5]