How to use transformed parameters in pytorch model (Module)?

adamr: <p><em>(This question is also posted on <a href="https://datascience.stackexchange.com/q/80352/60226" rel="nofollow noopener">https://datascience.stackexchange.com/q/80352/60226</a>. With permission, I will propagate here any answers I would get there).</em></p>
<p>Sometimes we might want to limit the domain of the parameter from the whole real, to e.g. positive reals. It may be accomplished by using a transform like</p>
<p><code>par -&gt; par**2 + epsilon</code></p>
<p>where <code>epsilon</code> is a constant governing how close to zero the value can land during the training process.</p>
<p>(Similarly, one can use <code>torch.sigmoid</code> function to the domain of the parameter from both ends.)</p>
<p>The problem is that from the code logic, the concept of transformed parameter belongs to a definition (initialization) of that parameter, whereas the actual function that does the transform must be called in the <code>forward</code> overload, otherwise we will get</p>
<pre><code class="lang-auto">RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling backward the first time.
</code></pre>
<p>when calling the <code>forward</code> for the second time. E.g.</p>
<pre><code class="lang-python">class GenericModule(torch.nn.Module):
    def __init__(self, weight):
        super(GenericModule, self).__init__()
        self.weight = weight
            
    def forward(self,X):
        return X * self.weight

class SpecificModule1(torch.nn.Module):
    def __init__(self):
        super(SpecificModule1, self).__init__()
        epsilon = 0.001
        self.specific_coef = torch.nn.Parameter(torch.Tensor([0.]))
        self.worker = GenericModule(self.specific_coef**2 + epsilon)
    def forward(self,X):
        return self.worker(X)
</code></pre>
<pre><code class="lang-auto">&gt;&gt;&gt; model = SpecificModule1()
&gt;&gt;&gt; list(model.named_parameters())
[('specific_coef',
  Parameter containing:
  tensor([0.], requires_grad=True))]
&gt;&gt;&gt; X = 1.
&gt;&gt;&gt; optim = torch.optim.RMSprop(model.parameters())
&gt;&gt;&gt; loss=model(X)
&gt;&gt;&gt; loss.backward()
&gt;&gt;&gt; loss=model(X)
&gt;&gt;&gt; loss.backward()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-114-52a0569421b1&gt; in &lt;module&gt;
----&gt; 1 loss.backward()

~/venv/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    183                 products. Defaults to ``False``.
    184         """
--&gt; 185         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    186 
    187     def register_hook(self, hook):

~/venv/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
    123         retain_graph = create_graph
    124 
--&gt; 125     Variable._execution_engine.run_backward(
    126         tensors, grad_tensors, retain_graph, create_graph,
    127         allow_unreachable=True)  # allow_unreachable flag

RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling backward the first time.

</code></pre>
<p>The usual way of dealing with such a situation is to hardwire the transformation into the <code>forward</code> function of <code>GenericModule</code>. Unfortunately, this has two problems.</p>
<ol>
<li>Hardwiring the transform into the child’s <code>forward</code> function is against DRY (Don’t Repeat Yourself) principle, because you may have to additionally spawn <code>GenericModuleWithRightBoundedPar</code> that will not be able to reuse the code in the <code>GenericModule</code>.</li>
<li>This will ban you from sharing (or just forwarding of) this transformed parameter with another module, child of <code>GenericModule</code>, unless you pass the transformed parameter in the <code>forward</code> method in <code>GenericModule</code> to its child - which is very awkward indeed because the child must be aware beforehand, that the parameter it uses may be transformed, so it accepts it as a parameter in its <code>forward</code> method rather than initialize it in the constructor following typical design patterns of writing Modules.</li>
</ol>
<h3>The questions:</h3>
<ol>
<li>Did I miss some simple solution of this problem?</li>
<li>Is this problem already solved by someone, so I can just import his library that derives his own classes from <code>torch.nn.Parameter</code> and <code>torch.nn.Module</code>, which supports transformed parameters?</li>
</ol>