code: def _compute_jacobian_wrt_params_with_sample_wise_trick(
      model: Module,
      inputs: Tuple[Any, ...],
      labels: Optional[Tensor] = None,
      loss_fn: Optional[Union[Module, Callable]] = None,
      reduction_type: Optional[str] = "sum",
      layer_modules: List[Module] = None,
  ) -> Tuple[Any, ...]:
      r"""
      Computes the Jacobian of a batch of test examples given a model, and optional
      loss function and target labels. This method uses sample-wise gradients per
      batch trick to fully vectorize the Jacobian calculation. Currently, only
      linear and conv2d layers are supported.
  
      User must `add_hooks(model)` before calling this function.
  
      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='sum'` or
                  `reduction='mean'`.
          reduction_type (str, optional): The type of reduction applied. If a loss_fn is
                  passed, this should match `loss_fn.reduction`. Else if gradients are
                  being computed on direct model outputs (scores), then 'sum' should be
                  used.
                  Defaults to 'sum'.
          layer_modules (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):
          inputs = tuple(inp.clone() for inp in inputs)
          apply_gradient_requirements(inputs)
          sample_grad_wrapper = SampleGradientWrapper(
              model, layer_modules
          )
          try:
              sample_grad_wrapper.add_hooks()
  
              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)
                  # TODO: allow loss_fn to be Callable
                  if (
                      isinstance(loss_fn, Module) or callable(loss_fn)
                  ) and hasattr(loss_fn, "reduction"):
                      reduction = loss_fn.reduction  # type: ignore
                      msg0 = "Please ensure that loss_fn.reduction is set to `sum` or `mean`"
  
                      assert reduction != "none", msg0
                      msg1 = (
                          f"loss_fn.reduction ({reduction}) does not match"
                          f"reduction type ({reduction_type}). Please ensure they are"
                          " matching."
                      )
                      assert reduction == reduction_type, msg1
                  msg2 = (
                      "Please ensure custom loss function is applying either a "
                      "sum or mean reduction."
                  )
                  assert out.shape != loss.shape, msg2
  
                  if (
                      reduction_type != "sum"
                      and reduction_type != "mean"
                  ):
                      raise ValueError(
                          f"{reduction_type} is not a valid value for reduction_type. "
                          "Must be either 'sum' or 'mean'."
                      )
                  out = loss
  
              sample_grad_wrapper.compute_param_sample_gradients(
                  out, loss_mode=reduction_type
              )
              if layer_modules is not None:
                  layer_parameters = _extract_parameters_from_layers(
                      layer_modules
                  )
              grads = tuple(
                  param.sample_grad  # type: ignore
                  for param in (
                      model.parameters()
                      if layer_modules is None
                      else layer_parameters
                  )
                  if hasattr(param, "sample_grad")
              )
          finally:
              sample_grad_wrapper.remove_hooks()
  
          return 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 24 for calling a global function with example usage: _compute_jacobian_wrt_params_with_sample_wise_trick(model, inputs, labels, loss_fn, reduction_type, layer_modules) and returns: grads