Question regarding efficiency of jax.grad

[Rasmuskh]

Hi,

Looking at the source code I can see that jax.grad calls jax.value_and_grad and just throws the value away. This seems a bit wasteful.

Example:
Say I want to compute $\frac{\partial}{\partial u}( u^T A v)$ , where u and v are vectors and A is a matrix.
Then it looks like jax.grad will call jax.value_and_grad and compute both the gradient $Av$ and the value $u^T A v$, even though computing $Av$ is sufficient.

My question is if this is actually what happens in practice or if perhaps the jit compiler is able to optimize away the unnecessary value computation?

[jakevdp]

Thanks for the question!

Typically in JAX, performance-critical code is wrapped in jax.jit, which is passed to the compiler, and the compiler will do dead code elimination before execution. So in practice, throwing away this value during trace time does not matter.

[Rasmuskh]

Thanks!
That is neat :)