def METHOD_NAME(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
           has_aux: bool = False, **kwargs) -> Callable:
  """Sparse-aware version of :func:`jax.jacfwd`
  Arguments and return values are the same as :func:`jax.jacfwd`, 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_jacfwd_fun = jax.METHOD_NAME(fun, argnums=argnums, **kwargs)
  argnums = core.concrete_or_error(_ensure_index, argnums)
  @wraps(fun, docstr=raw_jacfwd_fun.__doc__, argnums=argnums)
  @api_boundary
  def jacfwd_fun(*args, **kwargs):
    fun_flat, argnums_flat, args_flat, postprocess_gradients = flatten_fun_for_sparse_ad(fun, argnums, args)
    out = jax.METHOD_NAME(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 jacfwd_fun