Hi @tom! Thanks for the quick response!
When it comes to actually write what’s in the methods of jvp and vjp are there any examples?
For example, here’s an example for f(x) = x^3 (How to customize the double backward? - #4 by Naruto-Sasuke 1) where you’d just need to define the backward as grad_out * 3 * x**2. Which is in general just the derivative of the output of the layer w.r.t its input, which you then multiple by the grad_out term to get the gradient of the loss w.r.t the input.
Also, I’ve never seen the ctx.mul before, so I assume that’s a new feature in master?
Thank you for the help, it’s greatly apprecitated!