code: def test_augment(self):
      K = self.matrix
      n = K.shape[-1]
      m = n // 2
      Kaa = K[:, 0:m, 0:m]
      Laa = torch.linalg.cholesky(Kaa)
      Kbb = K[:, m:, m:]
  
      error_msg = "One and only one of `Kba` or `Lba` must be provided."
      with self.assertRaisesRegex(ValueError, error_msg):
          augment_cholesky(Laa, Kbb)
  
      Kba = K[:, m:, 0:m]
      L_augmented = augment_cholesky(Laa, Kbb, Kba)
      L = torch.linalg.cholesky(K)
      self.assertAllClose(L_augmented, L)
  
      # with jitter
      jitter = 3e-2
      Laa = torch.linalg.cholesky(
          Kaa + jitter * torch.eye(m).unsqueeze(0)
      )
      L_augmented = augment_cholesky(Laa, Kbb, Kba, jitter=jitter)
      L = torch.linalg.cholesky(K + jitter * torch.eye(n).unsqueeze(0))
      self.assertAllClose(L_augmented, L)
  
file_path: botorch/test/utils/probability/test_linalg.py
imports: from __future__ import annotations
  from copy import deepcopy
  import torch
  from botorch.utils.probability.linalg import augment_cholesky, PivotedCholesky
  from botorch.utils.testing import BotorchTestCase
label_desc: Use image node_id 8 for calling the TestPivotedCholesky obj's underlying member method code with example usage: obj.test_augment() without return types