code: def _compute_jacobian_wrt_params(
      model: Module,
      inputs: Tuple[Any, ...],
      labels: Optional[Tensor] = None,
      loss_fn: Optional[Union[Module, Callable]] = None,
      layer_modules: List[Module] = None,
  ) -> Tuple[Tensor, ...]:
      r"""
      Computes the Jacobian of a batch of test examples given a model, and optional
      loss function and target labels. This method is equivalent to calculating the
      gradient for every individual example in the minibatch.
  
      Args:
          model (torch.nn.Module): The trainable model providing the forward pass
          inputs (tuple[Any, ...]): The minibatch for which the forward pass is computed.
                  It is unpacked before passing to `model`, so it must be a tuple.  The
                  individual elements of `inputs` can be anything.
          labels (Tensor, optional): Labels for input if computing a loss function.
          loss_fn (torch.nn.Module or Callable, optional): The loss function. If a library
                  defined loss function is provided, it would be expected to be a
                  torch.nn.Module. If a custom loss is provided, it can be either type,
                  but must behave as a library loss function would if `reduction='none'`.
          layer_modules (List[torch.nn.Module], optional): A list of PyTorch modules
                   w.r.t. which jacobian gradients are computed.
      Returns:
          grads (tuple[Tensor, ...]): Returns the Jacobian for the minibatch as a
                  tuple of gradients corresponding to the tuple of trainable parameters
                  returned by `model.parameters()`. Each object grads[i] references to the
                  gradients for the parameters in the i-th trainable layer of the model.
                  Each grads[i] object is a tensor with the gradients for the `inputs`
                  batch. For example, grads[i][j] would reference the gradients for the
                  parameters of the i-th layer, for the j-th member of the minibatch.
      """
      with torch.autograd.set_grad_enabled(True):
          out = model(*inputs)
          assert (
              out.dim() != 0
          ), "Please ensure model output has at least one dimension."
  
          if labels is not None and loss_fn is not None:
              loss = loss_fn(out, labels)
              if hasattr(loss_fn, "reduction"):
                  msg0 = (
                      "Please ensure loss_fn.reduction is set to `none`"
                  )
                  assert loss_fn.reduction == "none", msg0  # type: ignore
              else:
                  msg1 = (
                      "Loss function is applying a reduction. Please ensure "
                      f"Output shape: {out.shape} and Loss shape: {loss.shape} "
                      "are matching."
                  )
                  assert loss.dim() != 0, msg1
                  assert out.shape[0] == loss.shape[0], msg1
              out = loss
  
          if layer_modules is not None:
              layer_parameters = _extract_parameters_from_layers(
                  layer_modules
              )
          grads_list = [
              torch.autograd.grad(
                  outputs=out[i],
                  inputs=cast(
                      Union[Tensor, Sequence[Tensor]],
                      model.parameters()
                      if layer_modules is None
                      else layer_parameters,
                  ),
                  grad_outputs=torch.ones_like(out[i]),
                  retain_graph=True,
              )
              for i in range(out.shape[0])
          ]
          grads = tuple([torch.stack(x) for x in zip(*grads_list)])
  
          return tuple(grads)
  
imports: import threading
  import typing
  import warnings
  from collections import defaultdict
  from typing import Any, Callable, cast, Dict, List, Optional, Sequence, Tuple, Union
  import torch
  from captum._utils.common import _reduce_list, _run_forward, _sort_key_list, _verify_select_neuron
  from captum._utils.sample_gradient import SampleGradientWrapper
  from captum._utils.typing import Literal, ModuleOrModuleList, TargetType, TensorOrTupleOfTensorsGeneric
  from torch import device, Tensor
  from torch.nn import Module
label_desc: Use image node_id 23 for calling a global function with example usage: _compute_jacobian_wrt_params(model, inputs, labels, loss_fn, layer_modules) and returns: tuple