Best practices for debugging vjp blowups

[lengstrom]

Hi, I have a complicated function I am taking a vjp over. The resulting product ends up being quite large. To debug this, I would like to inspect the compute graph (with the individual partials) of the reverse-mode autodiff procedure and understand why this ended up being the case. What are the best practices for doing this in jax? Is there a toolkit for logging these kinds of quantities beyond defining custom_vjps for each operation of interest?

Thank you!

[dfm]

I'm not sure that I have a great answer to the broad question here, but there are some options. One place to start is to try printing the jaxpr of your reverse pass:

print(jax.make_jaxpr(jax.vjp(fun, *args)[1])(*ct))

Which you can also combine with jax.named_scope to get a little bit more metadata. For example:

Example

import jax
import jax.numpy as jnp

@jax.jit
def fun(x):
  return f1(x) * f2(x)

@jax.jit
@jax.named_scope("f1")
def f1(x):
  return jnp.sin(x)

@jax.jit
@jax.named_scope("f2")
def f2(x):
  return jnp.exp(0.5 * x)

jax.make_jaxpr(jax.vjp(fun, jnp.ones(5))[1])(jnp.ones(5))

Prints:

{ lambda a:f32[5] b:f32[5] c:f32[5]; d:f32[5]. let
    e:f32[5] = pjit[
      name=fun
      jaxpr={ lambda ; f:f32[5] g:f32[5] h:f32[5] i:f32[5]. let
          j:f32[5] = mul h i
          k:f32[5] = pjit[
            name=f2
            jaxpr={ lambda ; l:f32[5] m:f32[5]. let
                n:f32[5] = mul m l
                o:f32[5] = mul 0.5 n
              in (o,) }
          ] g j
          p:f32[5] = mul i g
          q:f32[5] = pjit[
            name=f1
            jaxpr={ lambda ; r:f32[5] s:f32[5]. let t:f32[5] = mul s r in (t,) }
          ] f p
          u:f32[5] = add_any k q
        in (u,) }
    ] a b c d
  in (e,) }

This doesn't really help with logging the specific values, but perhaps it can provide some help identifying the problematic parts of the computation.

In terms of actually logging or otherwise intercepting the reverse mode computations, I think that custom_vjp probably is your best bet, but you might be able to automate some of that by calling jax.vjp in your custom_vjp.

Perhaps you can see how far something like this gets you, but please feel free to post a more specific example function and more details about what you're looking for if it would be useful to discuss further.