Calculate derivative of function

Blaze: <p>Suppose I have a function <code>def f(x): return torch.sin(x)</code></p>
<p>I can find the derivative of the function at a point as folllows:</p>
<pre><code class="lang-auto">x = torch.tensor([2.], requires_grad=True)
y = f(x)
y.backward(retain_graph=True)
x.grad
</code></pre>
<p>However, this requires computing the value of f(x) in order to find it’s derivative, which should not be required</p>
<p>Other libraries such as jax provide <code>grad(f)</code> which returns the function cos(x)</p>
<p>I was wondering if there is any Pytorch equivalent of this function?</p>
<p>I tried using <code>y.grad_fn</code> in conjunction with <code>next_functions</code> in order to reconstruct f’(x), but this did not seem to give correct values</p>
<p>Any ideas?</p>

albanD: <p>Hi,</p>
<p>Pytorch is a bit particular in the sense that the definition of your function is discovered at the same time as it is evaluated at a point. So you cannot get gradient without evaluating the function at a given point. (This allows much more flexibility on what is allowed wrt control flows and inplace operations).</p>
<p>If you want static graphs that are differentiated symbolically (without evaluation), you can turn to TorchScript which is going to do that.</p>

Blaze: <p>Thanks, that sounds like exactly what I’m looking for! However, after looking online for a while, I’m still confused as to how exactly to go about solving my problem using TorchScript?</p>

albanD: <p>You want to be able to the <code>grad</code> function directly?</p>

Blaze: <p>Yes, if it’s possible</p>

zdevito: <p>The TorchScript API for differentiation is the same as the eager mode API, so if it is not expressible in Python, it won’t be expressible in TorchScript as well. It is still possible to make something that appears like a grad function, but internally it will always compute the forward function first:</p>
<pre><code class="lang-auto">def grad(f):
    def result(x):
       # make leaf variables out of the inputs
       x_ = x.detach().requires_grad_(True) 
       f(x_).backward()
      return x_.grad
  return result
</code></pre>
<p>We do not do the kind of whole program transformations that mathematica or jax does that would make the generation of backwards possible.</p>

randbit: <p><a class="mention" href="/u/alband">@albanD</a> What’s the current recommended way to do this?</p>
<p><a class="mention" href="/u/zdevito">@zdevito</a> Is there a way to preserve gradients through that? If so, what’s the right way to do it?</p>
<p>To illustrate my use case (meta-optimization):</p>
<pre><code class="lang-auto">import torch

def grad(f):
    def result(x):
        x_ = x.detach().requires_grad_(True) 
        f(x_).backward()
        return x_.grad
    return result

def f(x):
    return x**2

def grad_f(x):
    g1 = 2 * x
    g2 = grad(f)(x)
    assert (g1 == g2).all()
    print('g1: {!r}\ng2: {!r}'.format(g1, g2))
    return g1

def optimize(lr):
    x = torch.tensor(1.)
    for _ in range(10**2):
        x = x - lr * grad_f(x)
    return x

lr = torch.tensor(1e-3, requires_grad=True)
opt = torch.optim.SGD([lr], 1e-6, .9)
while True:
    loss = f(optimize(lr))
    print('{:6.4f} {:6.4f}'.format(loss.item(), lr.item()))
    opt.zero_grad()
    loss.backward()
    opt.step()
</code></pre>
<p>Note that, after the first step, <code>g1</code> is connected to the computation graph defined by <code>lr</code>, but <code>g2</code> is not:</p>
<pre><code class="lang-auto">g1: tensor(2.)
g2: tensor(2.)
g1: tensor(1.9960, grad_fn=&lt;MulBackward0&gt;)
g2: tensor(1.9960)
g1: tensor(1.9920, grad_fn=&lt;MulBackward0&gt;)
g2: tensor(1.9920)
g1: tensor(1.9880, grad_fn=&lt;MulBackward0&gt;)
g2: tensor(1.9880)
g1: tensor(1.9840, grad_fn=&lt;MulBackward0&gt;)
g2: tensor(1.9840)
</code></pre>

albanD: <p>Hi,</p>
<p>I don’t think there is any specific API for this. Just compute the gradient the same way as you would do otherwise.<br>
Note that while jax provides a closure for grad, it might still re-compute the forward every time to use it in some cases.</p>

Sam_Lerman: <p>Just checking, does Pytorch have any plans to implement such static derivative functions, akin to Jax?</p>