code: def grad(fun, argnums=0, has_aux=False, holomorphic=False):
      r"""Creates a function which evaluates the gradient of ``fun``.
  
      NOTE: You only need this in a very specific case that you want to take a
      gradient **inside** a :func:`transform`\ ed function and the function you are
      differentiating uses :func:`set_state`. For example:
  
      >>> class MyModule(hk.Module):
      ...   def __call__(self, x):
      ...     hk.set_state("last", x ** 2)
      ...     return x ** 2
  
      >>> def f(x):
      ...   m = MyModule()
      ...   g = hk.grad(m)(x)
      ...   return g
  
      >>> f = hk.transform_with_state(f)
      >>> x = jnp.array(2.)
      >>> params, state = jax.jit(f.init)(None, x)
      >>> print(state["my_module"]["last"])
      4.0
  
      Args:
        fun: Function to be differentiated. Its arguments at positions specified by
          ``argnums`` should be arrays, scalars, or standard Python containers. It
          should return a scalar (which includes arrays with shape ``()`` but not
          arrays with shape ``(1,)`` etc.)
        argnums: Optional, integer or tuple of integers. Specifies which positional
          argument(s) to differentiate with respect to (default 0).
        has_aux: Optional, bool. Indicates whether ``fun`` returns a pair where the
          first element is considered the output of the mathematical function to be
          differentiated and the second element is auxiliary data. Default False.
        holomorphic: Optional, bool. Indicates whether ``fun`` is promised to be
          holomorphic. Default False.
  
      Returns:
        A function with the same arguments as ``fun``, that evaluates the gradient
        of ``fun``. If `argnums` is an integer then the gradient has the same shape
        and type as the positional argument indicated by that integer. If argnums is
        a tuple of integers, the gradient is a tuple of values with the same shapes
        and types as the corresponding arguments. If ``has_aux`` is True then a pair
        of ``gradient, auxiliary_data`` is returned.
  
      For example:
  
      >>> grad_tanh = jax.grad(jax.numpy.tanh)
      >>> print(grad_tanh(0.2))
      0.96...
      """
      value_and_grad_fun = value_and_grad(
          fun, argnums=argnums, has_aux=has_aux, holomorphic=holomorphic
      )
  
      @functools.wraps(fun)
      def grad_fn(*args, **kwargs):
          value, grads = value_and_grad_fun(*args, **kwargs)
          if has_aux:
              value, aux = value
              return grads, aux
          else:
              return grads
  
      return grad_fn
  
imports: import collections
  import collections.abc
  from collections.abc import Mapping, MutableMapping
  import functools
  import inspect
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  import jax
  import jax.numpy
label_desc: Use image node_id 7 for calling a global function with example usage: grad(fun, argnums, has_aux, holomorphic) and returns: grad_fn, grads,  aux, grads