code: def setUp(self):
      super().setUp()
      n = 5
      with torch.random.fork_rng():
          torch.random.manual_seed(0)
  
          matrix = torch.randn(2, n, n)
          matrix = matrix @ matrix.transpose(-1, -2)
  
          diag = matrix.diagonal(dim1=-2, dim2=-1).sqrt()
          idiag = diag.reciprocal().unsqueeze(-1)
  
          piv_chol = PivotedCholesky(
              step=0,
              tril=(idiag * matrix * idiag.transpose(-2, -1)).tril(),
              perm=torch.arange(n)[None]
              .expand(len(matrix), n)
              .contiguous(),
              diag=diag.clone(),
          )
  
      self.diag = diag
      self.matrix = matrix
      self.piv_chol = piv_chol
  
      self.piv_chol.update_()
      self.piv_chol.pivot_(torch.tensor([2, 3]))
      self.piv_chol.update_()
  
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 1 for calling the TestPivotedCholesky obj's underlying member method code with example usage: obj.setUp() without return types