code: def _jvp_with_argnums(
      func: Callable,
      primals: Any,
      tangents: Any,
      argnums: Optional[argnums_t],
      *,
      strict: bool = False,
      has_aux: bool,
  ):
      # This is the same function as jvp but also accepts an argnums argument
      # Most args are the same as jvp except for the added argument
      # argnums (Optional[int or tuple[int]]): Optional, specifies the argument(s) to compute gradients with respect to.
      #         If None, computes the gradients with respect to all inputs (used for jvp). Default: None
      # Because of this, tangents must be of length argnums and matches up to the corresponding primal whose index is
      # given by argnums
      #
      # WARN: Users should NOT call this function directly and should just be calling jvp.
      # It is only separated so that inputs passed to jacfwd but not differentiated get the correct wrappers.
      #
      # NOTE: All error messages are produced as if jvp was being called, even if this was called by jacfwd
      #
      # Returns the same two elements as :func:`jvp` but the returned tuple, ``jvp_out``, only has JVPs with respect to
      # the primals given by argnums
      if not isinstance(primals, tuple):
          raise RuntimeError(
              f"{jvp_str}: Expected primals to be a tuple. "
              f"E.g. it should be valid to call f(*primals)."
          )
      diff_args = (
          primals
          if argnums is None
          else _slice_argnums(primals, argnums)
      )
      flat_primals, primals_spec = tree_flatten(diff_args)
      flat_tangents, tangents_spec = tree_flatten(tangents)
      if primals_spec != tangents_spec:
          raise RuntimeError(
              f"{jvp_str}: Expected primals and tangents to have the same python "
              f"structure. For example, if primals is a tuple of 3 tensors, "
              f"tangents also must be. Got primals with structure {primals_spec} "
              f"and tangents with structure {tangents_spec}"
          )
      assert_non_empty_list_of_tensors(flat_primals, jvp_str, "primals")
      assert_non_empty_list_of_tensors(
          flat_tangents, jvp_str, "tangents"
      )
  
      level = _jvp_increment_nesting()
      try:
          global JVP_NESTING
          JVP_NESTING += 1
          with fwAD._set_fwd_grad_enabled(True):
              ctx = fwAD.dual_level if JVP_NESTING == 1 else noop
              with ctx():
                  flat_duals = tuple(
                      fwAD.make_dual(p, t)
                      for p, t in zip(flat_primals, flat_tangents)
                  )
                  duals = tree_unflatten(flat_duals, primals_spec)
                  if argnums is not None:
                      primals = _wrap_all_tensors(primals, level)
                      duals = _replace_args(primals, duals, argnums)
                  result_duals = func(*duals)
                  if has_aux:
                      if not (
                          isinstance(result_duals, tuple)
                          and len(result_duals) == 2
                      ):
                          raise RuntimeError(
                              f"{jvp_str}: output of function f should be a tuple: (output, aux) "
                              "if has_aux is True"
                          )
                      result_duals, aux = result_duals
                      aux = _undo_create_differentiable(aux, level)
  
                  result_duals, spec = tree_flatten(result_duals)
                  assert_non_empty_tensor_output(result_duals, jvp_str)
  
                  primals_out, tangents_out = zip(
                      *[
                          safe_unpack_dual(dual, strict)
                          for dual in result_duals
                      ]
                  )
                  primals_out = tree_map(
                      partial(_undo_create_differentiable, level=level),
                      primals_out,
                  )
                  tangents_out = tree_map(
                      partial(_undo_create_differentiable, level=level),
                      tangents_out,
                  )
  
                  primals_out_unflatten = tree_unflatten(
                      primals_out, spec
                  )
                  tangents_out_unflatten = tree_unflatten(
                      tangents_out, spec
                  )
                  if has_aux:
                      return (
                          primals_out_unflatten,
                          tangents_out_unflatten,
                          aux,
                      )
  
                  return primals_out_unflatten, tangents_out_unflatten
      finally:
          _jvp_decrement_nesting()
          JVP_NESTING -= 1
  
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 30 for calling a global function with example usage: _jvp_with_argnums(func, primals, tangents, argnums) and returns: primals_out_unflatten,  tangents_out_unflatten, primals_out_unflatten,  tangents_out_unflatten,  aux