Way to surface auxiliary data in custom gradient?

[dieterichlawson]

I have a custom gradient implementation that computes the gradient of a fixed point operation, very similar to the one in the JAX docs. I would like to output the number of iterations of fixed_point used to compute the gradient, but I see no easy way to return auxiliary data from the custom gradient of a function and surface it when the function is used in composition with a bunch of others (e.g. inside a big model). Are there any sanctioned ways to do this?

[flxrcz]

I am having a similar problem. Did you find any good way to do this?

[jakevdp]

You can do this using a placeholder in the custom JVP rule, and has_aux at the call-site; for example:

import jax
import jax.numpy as jnp

@jax.custom_jvp
def function(x):
    num_iter = 100
    return 2 * x, 100

@function.defjvp
def _function_jvp(primals, tangents):
    x, = primals
    x_dot, = tangents

    return function(x), (2 * x_dot, 0)

value, grad = jax.value_and_grad(function, has_aux=True)(5.0)
print(value)
# (10.0, 100)
print(grad)
# 2.0

[flxrcz]

Thanks for the quick answer,
is this also possible for a backwards pass in a vjp and composite gradients?
My current problem is very similar to the implicit differentiation example in the JAX guide:

from jax import vjp

@partial(custom_vjp, nondiff_argnums=(0,))
def fixed_point(f, a, x_guess):
  def cond_fun(carry):
    x_prev, x = carry
    return jnp.abs(x_prev - x) > 1e-6

  def body_fun(carry):
    _, x = carry
    return x, f(a, x)

  _, x_star = while_loop(cond_fun, body_fun, (x_guess, f(a, x_guess)))
  num_iter = 100
  return x_star

def fixed_point_fwd(f, a, x_init):
  x_star = fixed_point(f, a, x_init)
  return x_star, (a, x_star)

def fixed_point_rev(f, res, x_star_bar):
  a, x_star = res
  _, vjp_a = vjp(lambda a: f(a, x_star), a)
  a_bar, = vjp_a(fixed_point(partial(rev_iter, f),
                             (a, x_star, x_star_bar),
                             x_star_bar))
  return a_bar, jnp.zeros_like(x_star)
  
def rev_iter(f, packed, u):
  a, x_star, x_star_bar = packed
  _, vjp_x = vjp(lambda x: f(a, x), x_star)
  return x_star_bar + vjp_x(u)[0]

fixed_point.defvjp(fixed_point_fwd, fixed_point_rev)

def fixed_point_update(a, x_guess):
  # do something
  return x_guess_new

def function(a):
  x_guess = 0
  x = fixed_point(fixed_point_update, a, x_guess)
  # do something with x
  return x_transformed

grad(function)(1) # <- extract num_iter from the fixed point iteration, ideally for both the forward and the backward pass

So I want to differentiate my fixed point iteration with respect to some parameters and get info on how the forward fixed point iteration behaved and how the backwards fixed point iteration behaved. Is something like this possible with the has_aux keyword? I have already been able to extract some information via debugging, but I would like to do it more systematically.

Thanks in advance!