code: def grad(f, argnums=0, has_aux=False):
      """Returns the gradient function for `f`."""
      value_and_grad_f = value_and_grad(
          f, argnums=argnums, has_aux=has_aux
      )
  
      @functools.wraps(f)
      def wrapper(*args, **kwargs):
          if has_aux:
              (_, aux), g = value_and_grad_f(*args, **kwargs)
              return g, aux
          else:
              _, g = value_and_grad_f(*args, **kwargs)
              return g
  
      return wrapper
  
file_path: sonnet/sonnet/src/functional/jax.py
imports: import functools
  from sonnet.src.functional import utils
  import tensorflow
  import tree
label_desc: Use image node_id 4 for calling a global function with example usage: grad(f, argnums, has_aux) and returns: wrapper, g,  aux, g