Gradient calculation for numerically inverted function

username: bianDe
cooked: <p>You have two tensors x,z. Given z, say you have the analytical formula for  x = x(z) as a bijective function. So the gradient during training can be calculated. Now suppose you want to calculate z(x), when given x. This inverse function can only be calculated numerically (eg. bisection search) and to do this you would have to detach the result of the inversion. Obviously during this, it won’t be possible to calculate the gradient dz/dx with automatic differentiation. It is however possible to calculate the gradient as (dx/dz)^(-1). Theoretically, as well as with autograd.grad.</p>
<p>I am facing the following questions:</p>
<ul>
<li>Is there a way to reconnect the differentiation graph after .detach() for the bisection search ?</li>
<li>Is there a way to supply the correct gradients (dx/dz)^(-1) for the detached connection ?</li>
</ul>
<p>Any advice is welcome <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> Cheers</p>

username: tom
cooked: <p>Hi,</p>
<aside class="quote no-group" data-username="bianDe" data-post="1" data-topic="119276">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/5f8ce5/48.png" class="avatar"> bianDe:</div>
<blockquote>
<p>Is there a way to reconnect the differentiation graph after .detach() for the bisection search ?</p>
</blockquote>
</aside>
<p>Yes, you would want to write a custom <a href="https://pytorch.org/docs/stable/notes/extending.html" rel="noopener nofollow ugc">autograd.Function</a>. Note that you don’t need to detach the tensor as the forward is implicitly run</p>
<aside class="quote no-group" data-username="bianDe" data-post="1" data-topic="119276">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/b/5f8ce5/48.png" class="avatar"> bianDe:</div>
<blockquote>
<p>Is there a way to supply the correct gradients (dx/dz)^(-1) for the detached connection ?</p>
</blockquote>
</aside>
<p>Yes, you can use torch.autograd.grad to get dx/dz after computing x with a gradient-requiring copy of z in an torch.enable_grad() block . I would probably do this in the backward.</p>
<p>Best regards</p>
<p>Thomas</p>

username: bianDe
cooked: <p>Thanks a lot for your answer!<br>
In case there is interest. This is the solution I got at that time:</p>
<pre><code class="lang-auto">from torch.autograd import grad,Function

class MyFunctionInverse(Function):

    @staticmethod
    def forward(ctx, x, par_1):

        # get inverse
        func_partial = partial(myfunction, par_1=par_1)
        z = bisection_invert(func_partial, x)

        # save for backward
        ctx.save_for_backward(z, par_1)

        return z

    @staticmethod
    def backward(ctx, grad_z):

        with torch.enable_grad():

            z, par_1 = map(
                lambda t: t.detach().clone().requires_grad_(), ctx.saved_tensors)

            x = myfunction(z, par_1)

            grad_x_inverse, = grad(x, z, torch.ones_like(x), retain_graph=True)
            grad_x = grad_z*grad_x_inverse.pow(-1)

            # minus due to parameter dependency
            grad_par_1 = grad(
                x, [par_1], - grad_x)

        return grad_x, grad_par_1
</code></pre>

username: Azadeh_Karimisefat
cooked: <p>Have you set all your parameters of the model in [par_1]? such as [a,b,c,d]?</p>

username: bianDe
cooked: <p>Hi <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"> par_1 is a parameter of the function. Like eg. the slope in a linear function: f(x;a) = a*x. You can use several parameters.</p>