code: def jacfwd(
      func: Callable,
      argnums: argnums_t = 0,
      has_aux: bool = False,
      *,
      randomness: str = "error",
  ):
      """
      Computes the Jacobian of ``func`` with respect to the arg(s) at index
      ``argnum`` using forward-mode autodiff
  
      Args:
          func (function): A Python function that takes one or more arguments,
              one of which must be a Tensor, and returns one or more Tensors
          argnums (int or Tuple[int]): Optional, integer or tuple of integers,
              saying which arguments to get the Jacobian with respect to.
              Default: 0.
          has_aux (bool): Flag indicating that ``func`` returns a
              ``(output, aux)`` tuple where the first element is the output of
              the function to be differentiated and the second element is
              auxiliary objects that will not be differentiated.
              Default: False.
          randomness(str): Flag indicating what type of randomness to use.
              See :func:`vmap` for more detail. Allowed: "different", "same", "error".
              Default: "error"
  
      Returns:
          Returns a function that takes in the same inputs as ``func`` and
          returns the Jacobian of ``func`` with respect to the arg(s) at
          ``argnums``. If ``has_aux is True``, then the returned function
          instead returns a ``(jacobian, aux)`` tuple where ``jacobian``
          is the Jacobian and ``aux`` is auxiliary objects returned by ``func``.
  
      .. note::
          You may see this API error out with "forward-mode AD not implemented
          for operator X". If so, please file a bug report and we will prioritize it.
          An alternative is to use :func:`jacrev`, which has better operator coverage.
  
      A basic usage with a pointwise, unary operation will give a diagonal array
      as the Jacobian
  
          >>> from torch.func import jacfwd
          >>> x = torch.randn(5)
          >>> jacobian = jacfwd(torch.sin)(x)
          >>> expected = torch.diag(torch.cos(x))
          >>> assert torch.allclose(jacobian, expected)
  
      :func:`jacfwd` can be composed with vmap to produce batched
      Jacobians:
  
          >>> from torch.func import jacfwd, vmap
          >>> x = torch.randn(64, 5)
          >>> jacobian = vmap(jacfwd(torch.sin))(x)
          >>> assert jacobian.shape == (64, 5, 5)
  
      If you would like to compute the output of the function as well as the
      jacobian of the function, use the ``has_aux`` flag to return the output
      as an auxiliary object:
  
          >>> from torch.func import jacfwd
          >>> x = torch.randn(5)
          >>>
          >>> def f(x):
          >>>   return x.sin()
          >>>
          >>> def g(x):
          >>>   result = f(x)
          >>>   return result, result
          >>>
          >>> jacobian_f, f_x = jacfwd(g, has_aux=True)(x)
          >>> assert torch.allclose(f_x, f(x))
  
      Additionally, :func:`jacrev` can be composed with itself or :func:`jacrev`
      to produce Hessians
  
          >>> from torch.func import jacfwd, jacrev
          >>> def f(x):
          >>>   return x.sin().sum()
          >>>
          >>> x = torch.randn(5)
          >>> hessian = jacfwd(jacrev(f))(x)
          >>> assert torch.allclose(hessian, torch.diag(-x.sin()))
  
      By default, :func:`jacfwd` computes the Jacobian with respect to the first
      input. However, it can compute the Jacboian with respect to a different
      argument by using ``argnums``:
  
          >>> from torch.func import jacfwd
          >>> def f(x, y):
          >>>   return x + y ** 2
          >>>
          >>> x, y = torch.randn(5), torch.randn(5)
          >>> jacobian = jacfwd(f, argnums=1)(x, y)
          >>> expected = torch.diag(2 * y)
          >>> assert torch.allclose(jacobian, expected)
  
      Additionally, passing a tuple to ``argnums`` will compute the Jacobian
      with respect to multiple arguments
  
          >>> from torch.func import jacfwd
          >>> def f(x, y):
          >>>   return x + y ** 2
          >>>
          >>> x, y = torch.randn(5), torch.randn(5)
          >>> jacobian = jacfwd(f, argnums=(0, 1))(x, y)
          >>> expectedX = torch.diag(torch.ones_like(x))
          >>> expectedY = torch.diag(2 * y)
          >>> assert torch.allclose(jacobian[0], expectedX)
          >>> assert torch.allclose(jacobian[1], expectedY)
  
      """
  
      @wraps(func)
      def wrapper_fn(*args):
          error_if_complex("jacfwd", args, is_input=True)
          primals = (
              args if argnums is None else _slice_argnums(args, argnums)
          )
          flat_primals, primals_spec = tree_flatten(primals)
          flat_primals_numels = tuple(p.numel() for p in flat_primals)
          flat_basis = _construct_standard_basis_for(
              flat_primals, flat_primals_numels
          )
          basis = tree_unflatten(flat_basis, primals_spec)
  
          def push_jvp(basis):
              output = _jvp_with_argnums(
                  func, args, basis, argnums=argnums, has_aux=has_aux
              )
              # output[0] is the output of `func(*args)`
              error_if_complex("jacfwd", output[0], is_input=False)
              if has_aux:
                  _, jvp_out, aux = output
                  return jvp_out, aux
              _, jvp_out = output
              return jvp_out
  
          results = vmap(push_jvp, randomness=randomness)(basis)
          if has_aux:
              results, aux = results
              # aux is in the standard basis format, e.g. NxN matrix
              # We need to fetch the first element as original `func` output
              flat_aux, aux_spec = tree_flatten(aux)
              flat_aux = [value[0] for value in flat_aux]
              aux = tree_unflatten(flat_aux, aux_spec)
  
          jac_outs, spec = tree_flatten(results)
          # Most probably below output check can never raise an error
          # as jvp should test the output before
          # assert_non_empty_output(jac_outs, 'jacfwd(f, ...)(*args)')
  
          jac_outs_ins = tuple(
              tuple(
                  safe_unflatten(jac_out_in, -1, primal.shape)
                  for primal, jac_out_in in zip(
                      flat_primals,
                      jac_out.movedim(0, -1).split(
                          flat_primals_numels, dim=-1
                      ),
                  )
              )
              for jac_out in jac_outs
          )
          jac_outs_ins = tuple(
              tree_unflatten(jac_ins, primals_spec)
              for jac_ins in jac_outs_ins
          )
  
          if isinstance(argnums, int):
              jac_outs_ins = tuple(
                  jac_ins[0] for jac_ins in jac_outs_ins
              )
          if has_aux:
              return tree_unflatten(jac_outs_ins, spec), aux
          return tree_unflatten(jac_outs_ins, spec)
  
      return wrapper_fn
  
imports: from typing import Callable, Union, Tuple, List, Any, Optional
  import torch
  from functools import partial, wraps
  import contextlib
  from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map, tree_map_only, tree_map_, treespec_pprint
  from torch.utils import _pytree
  from torch.fx.experimental import const_fold
  from torch.fx.experimental.proxy_tensor import make_fx
  import torch.autograd.forward_ad
  from torch._subclasses.functional_tensor import FunctionalTensor
  from .vmap import doesnt_support_saved_tensors_hooks, get_chunk_sizes
  from .apis import vmap
  from torch._C._functorch import _wrap_for_grad, _unwrap_for_grad, _grad_increment_nesting, _grad_decrement_nesting, _jvp_increment_nesting, _jvp_decrement_nesting, _wrap_functional_tensor, _unwrap_functional_tensor, _func_decrement_nesting, _func_increment_nesting, _assert_wrapped_functional, _propagate_functional_input_mutation, set_inplace_requires_grad_allowed, get_inplace_requires_grad_allowed
  from torch._functorch.utils import exposed_in, argnums_t
label_desc: Use image node_id 32 for calling a global function with example usage: jacfwd(func, argnums, has_aux) and returns: wrapper_fn, tree_unflatten, jvp_out, tree_unflatten,  aux, jvp_out,  aux