custom_vjp with variable nondiff_argnums

[sashaDoubov]

Hi, I'm porting some pytorch code over to jax, and am having some trouble understanding how to implement similar functionality in jax.

In the pytorch code, the authors have a custom forward and backwards function. In the custom backwards function, they optionally return the gradient wrt a variable depending on whether the input requires a gradient, otherwise they return None.

@torch.cuda.amp.custom_fwd
def f(ctx, x, y, z):
    # performs some computation
    return out
@torch.cuda.amp.custom_bwd
def backward(ctx, dout):
    ...
    dx, dy, dz = None, None, None
    if ctx.needs_input_grad[2]:
        dz = ...  # perform expensive computation
    if ctx.needs_input_grad[1]:
        dy = ...
    if ctx.needs_input_grad[0]:
        dx = ...
    return dx, dy, dz

However, in jax, from reading the docs, it seems to me that I would have to either:

1. Always compute dx, dy, dz fully
2. Use @partial(custom_vjp,nondiff_argnums=(0,)), but since I don't know ahead of time which variable doesn't require a gradient, I would somehow need to use the correct function depending on which variables I need the gradient of.

Is there a way to implement similar logic in jax? Or can this be avoided somehow?

[jakevdp]

In JAX, I believe the best way to do this would be to use f.defjvps, and pass a function per derivative rule. For example:

import jax
import jax.numpy as jnp

@jax.custom_jvp
def f(x, y, z):
    return x * jnp.sin(y) * jnp.cos(z)

f.defjvps(lambda x_dot, primal_out, x, y, z: x_dot * jnp.sin(y) * jnp.cos(z),
          lambda y_dot, primal_out, x, y, z: y_dot * x * jnp.cos(y) * jnp.cos(z),
          lambda z_dot, primal_out, x, y, z: -z_dot * x * jnp.sin(y) * jnp.sin(z))

print(jax.grad(f, argnums=0)(1.0, 1.0, 1.0))
# 0.45464867

You can see some examples of this in the Custom Derivative Rules docs.

[sashaDoubov]

Thank you, this is helpful! In my case, I actually need to perform a shared expensive computation first, in pytorch:

@torch.cuda.amp.custom_bwd
def backward(ctx, dout):
    ...
    dx, dy, dz = None, None, None
    out = ... # perform expensive initial computation
    if ctx.needs_input_grad[2]:
        dz = out * ...  # perform expensive computation
    if ctx.needs_input_grad[1]:
        dy = out * ...
    if ctx.needs_input_grad[0]:
        dx = out * ...
    return dx, dy, dz

In the jax case, will each of the associated functions in f.defjvps need me to recompute out = ..., or is there a way to make sure that expensive shared computation is performed first?

[jakevdp]

I see - in this case I think the best you can do is compute them in one function, something like this:

@f.defjvp
def _f_jvp(primals, tangents):
  x, y, z = primals
  x_dot, y_dot, z_dot = tangents
  primal_out = f(x, y, z)

  # do some common computation here
  sin_y = jnp.sin(y)
  cos_z = jnp.cos(z)

  df_dx = sin_y * cos_z
  df_dy = x * jnp.cos(y) * cos_z
  df_dz = -x * sin_y * jnp.sin(z)
  tangent_out = x_dot * df_dx + y_dot * df_dy + z_dot * df_dz
  return primal_out, tangent_out

Then you can rely on XLA's dead-code elimination within JIT to eliminate parts of the computation that are unnecessary. I may be wrong, but I don't think there's any other mechanism within JAX's autodiff machinery to define things at the granularity that you have in mind.

[jakevdp]

You can see this dead code by printing the jaxpr; for example:

print(jax.make_jaxpr(jax.grad(f, argnums=0))(1.0, 1.0, 1.0))

{ lambda ; a:f32[] b:f32[] c:f32[]. let
    _:f32[] = custom_jvp_call[
      call_jaxpr={ lambda ; d:f32[] e:f32[] f:f32[]. let
          g:f32[] = sin e
          h:f32[] = mul d g
          i:f32[] = cos f
          j:f32[] = mul h i
        in (j,) }
      jvp_jaxpr_thunk=<function _memoize.<locals>.memoized at 0x7fd30b919a70>
      num_consts=0
    ] a b c
    k:f32[] = sin b
    l:f32[] = cos c
    m:f32[] = mul k l
    n:f32[] = cos b
    o:f32[] = mul a n
    p:f32[] = mul o l
    q:f32[] = neg a
    r:f32[] = mul q k
    s:f32[] = sin c
    t:f32[] = mul r s
    _:f32[] = mul 0.0 p
    _:f32[] = mul 0.0 t
    u:f32[] = mul 1.0 m
  in (u,) }

The two statements that look like _:f32[] = mul 0.0 p are the results of the unused derivatives, and during compilation they should be trimmed from the computation by XLA.

[jakevdp]

Also, the whole "common computation" thing should be taken care of during compilation as well (repeated identical operations will be computed only once) so I'd suggest defining your jvp functions in the way that's easiest to code rather than worrying about the details of intermediate results and unnecessary cycles: the compiler should handle all of those considerations for you.

[sashaDoubov]

Perfect, this answers my question! I will just explicitly write out the jvps as you mentioned and rely on the compiler more heavily. Thank you!