code: def value_and_grad(fn, args):
      """Given `fn: (args) -> out, extra`, returns `dout/dargs`."""
      output, vjp_fn, extra = jax.vjp(fn, args, has_aux=True)
      grad = vjp_fn(jnp.ones_like(output))[0]
      return output, extra, grad
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  from jax import tree_util
  import jax.numpy
label_desc: Use image node_id 7 for calling a global function with example usage: value_and_grad(fn, args) and returns: output,  extra,  grad