How to implement a new function and its gradient in PyTorch?

arachnidman: <p>I want to implement a function that cannot be easily expressed with functions already implemented in PyTorch. However, I know how to compute this function and its gradients numerically. How could I implement this function as something that plays nicely with other PyTorch functions?</p>

Matias_Vasquez: <p>I think that this is too general to answer.<br>
The general answer would be to try and make a workaround using Pytorch functions.<br>
If this is not possible then maybe you can look into defining your own <code>autograd function</code>. <a href="https://pytorch.org/tutorials/beginner/examples_autograd/polynomial_custom_function.html">Here</a> is an example on how to do this.</p>
<p><a href="https://pytorch.org/docs/stable/notes/extending.html">Here</a> is a little more information on this.</p>
<p>Maybe you can post the equation or explain how this function should behave and maybe someone can come up with a solution or at least help you get in the right direction.</p>

arachnidman: <p>I am curious: Is there a way to see the gradients for the already implemented functions?</p>

AlphaBetaGamma96: <p>The backward formulae for basic PyTorch ops are defined in C++ source code which you can view on github, I don’t believe it’s accessible within the python API. If you want to define a new gradient do what <a class="mention" href="/u/matias_vasquez">@Matias_Vasquez</a> already stated and use <code>torch.autograd.Function</code> to define a custom function.</p>

Matias_Vasquez: <p>Additionally to what <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a> said, you could plot the gradients to visually inspect them by doing something like this.</p>
<pre data-code-wrap="python"><code class="lang-python">def func(x):
    return torch.nn.ReLU()(4*x**2 - 2)-1
    
x = torch.arange(-2, 2.01, 0.01, requires_grad=True)

y = func(x)
y.backward(gradient=torch.ones_like(x))

y_prime = x.grad

plt.plot(x.detach().numpy(), y.detach().numpy(), label="y")
plt.plot(x.detach().numpy(), y_prime.numpy(), label="y'")
plt.legend()
plt.show()
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/d/e/de4e22ecbca25fd91feb41f3606e10bb6ce802d2.png" alt="image" data-base62-sha1="vIBp6EMnYIDy2z5CYqfsKFokbCO" width="377" height="248"></p>
<p>But this might only make sense for some functions. Others are just going to seem random noise.</p>
<pre data-code-wrap="python"><code class="lang-python"># Using the same code as above, only changing the function
def func(x):
    return torch.nn.Linear(x.shape[-1], x.shape[-1])(x)
</code></pre>
<p>The result might look something like this ↓<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/e/c/ec8f68938f459c9695c12ddf2a03f1b1131edca0.png" alt="image" data-base62-sha1="xKHVVxtg8KgD5y21dH77nMcQnZu" width="370" height="248"></p>

AlphaBetaGamma96: <p>If you could share more about your use case, it would help clarify the situation.</p>
<p>Using <code>torch.autograd.Function</code> is good to define custom functions, but it can get messy very quickly as you have to manually define the backward component of your function.</p>