def _f(params, *args):
    """The gradient function to be returned."""
    with tf.GradientTape() as g:
      g.watch(tf.nest.flatten(params))
      outputs = f(params, *args)
      if has_aux:
        np_loss, aux = outputs
      else:
        np_loss = outputs
      check_loss_shape(np_loss)
      tf_grads = g.gradient(np_loss, params)
      if has_aux:
        res = (tf_grads, aux)
      else:
        res = tf_grads
      return _tf_to_np(res)
The gradient function to be returned.
trax/tf_numpy/extensions/extensions.py