I would like to implement a layer with weight sharing, that performs the matrix multiplication M @ x, with
M = [
[a, b]
[c, a]
]
where a, b, and c are the unique weights.
How can I implement this layer?
I tried instantiating M as a nn.Linear module, then doing M.weight[1][1] = M.weight[0][0]. While this works for weight sharing, the forward pass complexity is still the same as a non-weight-sharing layer; same for the model size.
Any help would be greatly appreciated, thanks.