code: def jacrev(
      fun: Callable,
      argnums: int | Sequence[int] = 0,
      has_aux: bool = False,
      **kwargs,
  ) -> Callable:
      """Sparse-aware version of :func:`jax.jacrev`
  
      Arguments and return values are the same as :func:`jax.jacrev`, but when taking
      the gradient with respect to a :class:`jax.experimental.sparse` array, the
      gradient is computed in the subspace defined by the array's sparsity pattern.
      Currently this is only implemented for dense outputs.
      """
      raw_jacrev_fun = jax.jacrev(fun, argnums=argnums, **kwargs)
      argnums = core.concrete_or_error(_ensure_index, argnums)
  
      @wraps(fun, docstr=raw_jacrev_fun.__doc__, argnums=argnums)
      @api_boundary
      def jacrev_fun(*args, **kwargs):
          (
              fun_flat,
              argnums_flat,
              args_flat,
              postprocess_gradients,
          ) = flatten_fun_for_sparse_ad(fun, argnums, args)
          out = jax.jacrev(
              fun_flat, argnums=argnums_flat, has_aux=has_aux, **kwargs
          )(*args_flat)
          if has_aux:
              return postprocess_gradients(out[0]), out[1]
          return postprocess_gradients(out)
  
      return jacrev_fun
  
file_path: jax/jax/experimental/sparse/ad.py
imports: from __future__ import annotations
  from collections.abc import Sequence
  import itertools
  from typing import Any, Callable, Union
  import jax
  from jax._src import core
  from jax import tree_util
  from jax._src.api_util import _ensure_index, _ensure_index_tuple
  from jax.util import safe_zip
  from jax._src.util import split_list, wraps
  from jax._src.traceback_util import api_boundary
  from jax.experimental.sparse._base import JAXSparse
label_desc: Use image node_id 5 for calling a global function with example usage: jacrev(fun, argnums, has_aux) and returns: jacrev_fun, postprocess_gradients, postprocess_gradients,  out