how to implement the vjp for the basic primitive

[mmmeee1111]

Hi, experts! I am learning to implement auto-diff for basic primitive, for example, numpy.real and numpy.imag.
Could you please point me to code or share the example on how to implement the forward and backward for these two operations? Thanks

[mmmeee1111]

@jakevdp could you help on this?

[jakevdp]

Here is where the autodiff rules are defined for the real_p and imag_p primitives in JAX: https://github.com/google/jax/blob/292deef6fda9f639fcecd9883a1112825a1eb54f/jax/_src/lax/lax.py#L1890-L1896

They use ad.deflinear2, which is a utility for registering jvp and transpose rules for linear primitives: https://github.com/google/jax/blob/292deef6fda9f639fcecd9883a1112825a1eb54f/jax/_src/interpreters/ad.py#L517-L519

For more general information regarding autodiff in JAX, some good resources from the docs are

• https://jax.readthedocs.io/en/latest/autodidax.html
• https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#how-it-s-made-two-foundational-autodiff-functions

And if you're not working with primitives, but rather want to define custom JVP / VJP rules for a function, there is some information here: https://jax.readthedocs.io/en/latest/_autosummary/jax.custom_jvp.html

Hope that helps!

[mmmeee1111]

@jakevdp Hi Jake, these are super helpful. I have one more specific question.

I am trying to understand the implementation. Could you point me to code that calculate the primal and residual value for the backward pass? for example, given the input x+yj, I think we should do v^T*(Derivative of (x+yj))? But I can't match the math with the implementation of the code.

[jakevdp]

The backward pass derivative is constructed from the forward-mode linearization (linear_jvp here) followed by the transpose rule, which is the second argument to the deflinear2 calls I linked to above. You can read about the math behind that decomposition here: https://arxiv.org/pdf/2105.09469.pdf