Deferred operations for advanced weight sharing

Dominik: <p>I develop differentiable simulations with PyTorch and I want to share weights (<code>torch.nn.Parameter</code>'s) between multiple components (<code>torch.nn.Module</code>'s). The crux is that one components should use the inverse (negative) weights of another.</p>
<p>Consider the following <code>Module</code> class:</p>
<pre><code>class Component(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.attr = torch.tensor(1.)

    def forward(self, x):
        return self.attr * x
</code></pre>
<p>Now if I have two such components <code>c1</code> and <code>c2</code> I could do “normal” weight sharing easily via</p>
<pre><code>c1.attr = c2.attr = torch.nn.Parameter(c1.attr)
</code></pre>
<p>However I need that <code>c2.attr == -c1.attr</code> at all times. If I do</p>
<pre><code>c1.attr = torch.nn.Parameter(c1.attr)
c2.attr = -c1.attr
</code></pre>
<p>then the operation <code>-c1.attr</code> creates a new node in the graph and doesn’t refer to the original parameter anymore. If I now use the two tensors <code>c1.attr</code> and <code>c2.attr</code> during multiple epochs the value of <code>c2.attr</code> won’t update since it is “stuck” at the initially captured value of <code>-c1.attr</code>.</p>
<p>What I do right now is update all these weight sharing relationships after each call to <code>optimizer.step()</code>. This however moves the logic of my model outside of its own definition and requires the user to not forget to call that update method every epoch, otherwise the weight sharing is ineffective.</p>
<p>Hence I would like to ask whether there exists a way for specifying some kind of deferred operations. That is operations that are only evaluated when used elsewhere in another computation. So for example:</p>
<pre><code>c2.attr = defer_mul(-1, c1.attr)
</code></pre>
<p>would compute <code>-1 * c1.attr</code> only when requested in another computation (and not cache any results). So <code>x * c2.attr</code> would result in <code>x * (-1 * c1.attr)</code> using the up-to-date value of <code>c1.attr</code> every time <code>c2.attr</code> is used.</p>
<p>Here is a more complete code example to illustrate the situation:</p>
<pre><code>import torch


class Component(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.attr = torch.tensor(1.)

    def forward(self, x):
        return self.attr * x


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.config = torch.tensor(-2.)
        self.c1 = Component()
        self.c2 = Component()

        self.c1.attr = torch.nn.Parameter(self.c1.attr)
        # self.c2.attr = -self.c1.attr  # stays at `attr = -1` forever
        self.update()  # update manually instead

    def forward(self):
        return self.c2(self.c1(self.config))

    def update(self):
        self.c2.attr = -self.c1.attr


model = Model()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)

expected = torch.tensor(32.)
for epoch in range(25):
    optimizer.zero_grad()
    value = model()
    loss = (value - expected) ** 2
    loss.backward()
    optimizer.step()
    model.update()  # need to update weight sharing relationships at every epoch

    print(f'c1 = {model.c1.attr}, c2 = {model.c2.attr}')</code></pre>