Hi,
I’m relatively new to Pytorch and am trying to play with a toy-demo to understand the autograd functionality better. I’m trying to create a custom nn.module which is similar to the Linear layer except the outgoing weights are tied (in a per-input way). E.g. suppose I have input x = [x1 x2 x3 x4] and 4 hidden units. I want the weight matrix to look like:
W = [ w1 w1 w1 w1
  w2 w2 w2 w2
  w3 w3 w3 w3
  w4 w4 w4 w4]

That is, I effectively only want to store 4 weights and 4 biases in this layer [w1 w2 w3 w4] and [b1 b2 b3 b4]. But during the computation of forward() and backward(), I’m hoping the stored weights can be expanded into the full matrix W and hiddens can be calculated as:
h = x * W
I’m in particular unsure how to implement the “storing [w1 w2 w3 w4] and expanding into full W matrix” part. I’m not sure if autograd will send the gradient information back through all 4x4=16 “effective weights” at computation time. I understand this model isn’t very useful but I’m doing it to understand how to manipulate the nn.modules at a finer level so I can do more useful things in the future.
Here is an attempt at a solution:

class tiedLinear(nn.Module):
    def __init__(self, in_features, bias=True):
        super(tiedLinear, self).__init__()
        self.in_features = in_features
        self.weight = Parameter(torch.Tensor(in_features))
        if bias:
            self.bias = Parameter(torch.Tensor(in_features))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(0))
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input):
        repeated_weight = None
        for i in range(0, self.in_features):
            row = self.repeat(self.weight[i], self.in_features)
            repeated_weight = repeated_weight.cat(repeated_weight, row, 0)

        return F.linear(input, self.repeated_weight, self.bias)

    def repeat(self, w, total):
        return torch.mul(torch.ones(total), w)