Define a custom variable pytorch tensor
I have a problem similar to this one:

There is an input vector of the form (n1, n2), and I want a model where f(n1, n2) = f(n2, n1). I want this model to be linear. For instance a 2x2 matrix that I want to learn. The 2x2 matrix has 4 weights which the training will try to learn. However, the constriction allows a clever reduction of the problem. 
f(nvec) = W*nvec, being nvec=(n1, n2).
So, let X = [[0, 1],[1, 0]] be a 2x2 matrix that flips vectors, then f(n2,n1) = W*X*(n1, n2) = f(n1, n2) = W (n1, n2).

Basically, the equality implies that W*X = W, or if W = [[w1, w2], [w3, w4]], then the equality implies W = [[w1, w1], [w3, w3]], i.e. the matrix has only 2 free parameters. How can I define a model like this one in pytorch?

Thanks.

If you have a special structure to your parameters, and you do not want to use the default parameters stored by the default layers, you can define a layer on your own and use nn.Parameter to store your parameters.  

For example,

class CustomLayer(nn.Module):  # layer must be derived from nn.Module class
  def __init__(self):  # you can have arguments here...
    super(CustomLayer, self).__init__()
    self.params = nn.Parameter(data=torch.rand((2, 1), dtype=torch.float), requires_grad=True)

  def forward(self, x):
    w = self.params.repeat(1, 2)  # create a 2x2 matrix from the two parameters
    y = torch.bmm(x, w)
    return y

