code: def batch_jacobian(
      f,
      *args,
      auto_unpack_single_arg=True,
      has_aux=False,
      name=None,
      **kwargs,
  ):
      """Computes batch Jacobian of `f(*args, **kwargs)` wrt to `args`, `kwargs`.
  
      Args:
        f: Python `callable`, returning a 2D `(batch, n)` shaped `Tensor`.
        *args: Arguments as in `f(*args, **kwargs)` and basis for Jacobian. Each
          element must be 2D `(batch, n)`-shaped argument `Tensor`(s). If multiple
          are provided, a tuple of jacobians are returned.
        auto_unpack_single_arg: Python `bool` which when `False` means the single
          arg case will not be interpreted as a list of arguments.
          Default value: `True`.
        has_aux: Whether `f(*args, **kwargs)` actually returns two outputs, the
          first being `y` and the second being an auxiliary output that does not get
          gradients computed.
        name: Python `str` name prefixed to ops created by this function.
          Default value: `None` (i.e., `'value_and_gradient'`).
        **kwargs: Named arguments as in `f(*args, **kwargs)` and basis for Jacobian.
          Each element must be 2D `(batch, n)`-shaped argument `Tensor`(s). If
          multiple are provided, a tuple of jacobians are returned.
  
      Returns:
        If `has_aux` is `False`:
          jacobian: A `(batch, n, n)` shaped `Tensor`, `dy/dx`, or a tuple thereof.
        otherwise:
          jacobian: A `(batch, n, n)` shaped `Tensor`, `dy/dx`, or a tuple thereof.
          aux: The auxiliary output of the function `y, aux = f(*args, **kwargs)`.
      """
      res = value_and_batch_jacobian(
          f,
          *args,
          auto_unpack_single_arg=auto_unpack_single_arg,
          name=name,
          has_aux=has_aux,
          **kwargs,
      )
      if has_aux:
          (_, aux), jacobian = res
          return jacobian, aux
      else:
          _, jacobian = res
          return jacobian
  
file_path: probability/tensorflow_probability/python/math/gradient.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow.python.util import tf_inspect
label_desc: Use image node_id 4 for calling a global function with example usage: batch_jacobian(f) and returns: jacobian,  aux, jacobian