code: def test_jacobian_loss_batch_vector_multilayer(self) -> None:
      model = BasicLinearModel_Multilayer(5, 2, 2)
      model.linear1.weight = nn.Parameter(
          torch.arange(0, 10).view(2, 5).float()
      )
      model.linear2.weight = nn.Parameter(
          torch.arange(0, 4).view(2, 2).float()
      )
  
      a = torch.stack((torch.ones(5), torch.ones(5) * 2))
      label = torch.Tensor([[33, 124], [69, 256]])
  
      loss_fn = nn.MSELoss(reduction="none")
      grads = _compute_jacobian_wrt_params(model, (a,), label, loss_fn)
      assertTensorAlmostEqual(
          self,
          grads[0][0],
          torch.stack(
              (
                  2 * (0 * (35 - 33) + 2 * (125 - 124)) * a[0],
                  2 * (1 * (35 - 33) + 3 * (125 - 124)) * a[0],
              )
          ),
      )
      assertTensorAlmostEqual(
          self,
          grads[1][0],
          torch.Tensor(
              [
                  [2 * (35 - 33) * 10, 2 * (35 - 33) * 35],
                  [2 * (125 - 124) * 10, 2 * (125 - 124) * 35],
              ]
          ),
      )
      assertTensorAlmostEqual(
          self,
          grads[0][1],
          torch.stack(
              (
                  2 * (0 * (70 - 69) + 2 * (250 - 256)) * a[1],
                  2 * (1 * (70 - 69) + 3 * (250 - 256)) * a[1],
              )
          ),
      )
      assertTensorAlmostEqual(
          self,
          grads[1][1],
          torch.Tensor(
              [
                  [2 * (70 - 69) * 10 * 2, 2 * (70 - 69) * 35 * 2],
                  [2 * (250 - 256) * 10 * 2, 2 * (250 - 256) * 35 * 2],
              ]
          ),
      )
  
      loss_fn = nn.MSELoss(reduction="sum")
      grads_h = _compute_jacobian_wrt_params_with_sample_wise_trick(
          model, (a,), label, loss_fn
      )
      assertTensorAlmostEqual(self, grads_h[0][0], grads[0][0])
      assertTensorAlmostEqual(self, grads_h[1][0], grads[1][0])
      assertTensorAlmostEqual(self, grads_h[0][1], grads[0][1])
      assertTensorAlmostEqual(self, grads_h[1][1], grads[1][1])
  
imports: import torch
  import torch.nn
  from captum._utils.gradient import _compute_jacobian_wrt_params, _compute_jacobian_wrt_params_with_sample_wise_trick
  from tests.helpers.basic import assertTensorAlmostEqual, BaseTest
  from tests.helpers.basic_models import BasicLinearModel2, BasicLinearModel_Multilayer
label_desc: Use image node_id 14 for calling the Test obj's underlying member method code with example usage: obj.test_jacobian_loss_batch_vector_multilayer() without return types