docstr: Test `augment`.
docstr_tok: ["Test","`","augment","`","."]
code: def test_augment(self):
      r"""Test `augment`."""
      with torch.random.fork_rng():
          torch.random.manual_seed(next(self.seed_generator))
  
          # Pick a set of subproblems at random
          index = torch.randint(
              low=0,
              high=len(self.sqrt_covariances),
              size=(),
              device=self.device,
          )
          sqrt_cov = self.sqrt_covariances[index]
          cov = sqrt_cov @ sqrt_cov.transpose(-2, -1)
          bounds = self.bounds[index]
  
          # Partially solve for `N`-dimensional integral
          N = cov.shape[-1]
          n = torch.randint(low=1, high=N - 2, size=())
          full = MVNXPB(cov, bounds=bounds)
          full.solve(num_steps=n)
  
          # Compare with solver run using a pre-computed `piv_chol`
          _perm = torch.arange(0, N, device=self.device)
          other = MVNXPB.build(
              step=0,
              perm=_perm.expand(*cov.shape[:-2], N).clone(),
              bounds=cov.diagonal(dim1=-2, dim2=-1)
              .rsqrt()
              .unsqueeze(-1)
              * bounds.clone(),
              piv_chol=full.piv_chol,
              plug_ins=full.plug_ins,
              log_prob=torch.zeros_like(full.log_prob),
          )
          other.solve(num_steps=n)
          self.assertTrue(full.perm.equal(other.perm))
          self.assertTrue(full.bounds.allclose(other.bounds))
          self.assertTrue(full.log_prob.allclose(other.log_prob))
  
          # Reorder terms according according to `full.perm`
          perm = full.perm.detach().clone()
          _cov = cov.gather(-2, perm.unsqueeze(-1).repeat(1, 1, N))
          _cov = _cov.gather(-1, perm.unsqueeze(-2).repeat(1, N, 1))
          _istd = _cov.diagonal(dim1=-2, dim2=-1).rsqrt()
          _bounds = bounds.gather(
              -2, perm.unsqueeze(-1).repeat(1, 1, 2)
          )
  
          # Solve for same `n`-dimensional integral as `full.solve(num_steps=n)`
          init = MVNXPB(_cov[..., :n, :n], _bounds[..., :n, :])
          init.solve()
  
          # Augment solver with adaptive pivoting disabled
          with patch.object(init.piv_chol, "diag", new=None):
              _corr = (
                  _istd[..., n:, None]
                  * _cov[..., n:, :]
                  * _istd[:, None, :]
              )
              temp = init.augment(
                  covariance_matrix=_corr[..., n:],
                  cross_covariance_matrix=_corr[..., :n],
                  bounds=_istd[..., n:, None] * _bounds[..., n:, :],
                  disable_pivoting=True,
              )
          self.assertTrue(temp.piv_chol.diag[..., :n].eq(1).all())
          self.assertEqual(temp.step, n)
          self.assertEqual(temp.piv_chol.step, N)
          self.assertTrue(
              temp.piv_chol.perm[..., n:].eq(_perm[n:]).all()
          )
          del temp
  
          # Augment solver again, this time with pivoting enabled
          augm = init.clone().augment(
              covariance_matrix=_cov[..., n:, n:],
              cross_covariance_matrix=_cov[..., n:, :n],
              bounds=_bounds[..., n:, :],
          )
  
          # Patch `perm` to account for different starting points
          augm_perm = augm.perm
          temp_perm = perm.gather(-1, augm_perm)
          self.assertTrue(augm_perm.equal(augm.piv_chol.perm))
          with patch.object(augm, "perm", new=temp_perm), patch.object(
              augm.piv_chol, "perm", new=temp_perm
          ):
              self.assertEqualMXNBPB(full, augm)
  
          # Run both solvers
          augm.piv_chol = full.piv_chol.clone()
          augm.piv_chol.perm = augm_perm.clone()
          full.solve()
          augm.solve()
  
          # Patch `perm` to account for different starting points
          augm_perm = augm.perm
          temp_perm = perm.gather(-1, augm_perm)
          self.assertTrue(augm_perm.equal(augm.piv_chol.perm))
          with patch.object(augm, "perm", new=temp_perm), patch.object(
              augm.piv_chol, "perm", new=temp_perm
          ):
              self.assertEqualMXNBPB(full, augm)
  
          # testing errors
          fake_init = deepcopy(init)
          fake_init.piv_chol.step = fake_init.perm.shape[-1] + 1
          error_msg = "Augmentation of incomplete solutions not implemented yet."
          with self.assertRaisesRegex(NotImplementedError, error_msg):
              augm = fake_init.augment(
                  covariance_matrix=_cov[..., n:, n:],
                  cross_covariance_matrix=_cov[..., n:, :n],
                  bounds=_bounds[..., n:, :],
              )
  
          # Testing that solver will try to recover if it encounters
          # a non-psd matrix, even if it ultimately fails in this case
          error_msg = "Matrix not positive definite after repeatedly adding jitter up to.*"
          with self.assertRaisesRegex(NotPSDError, error_msg):
              fake_cov = torch.ones_like(_cov[..., n:, n:])
              augm = init.augment(
                  covariance_matrix=fake_cov,
                  cross_covariance_matrix=_cov[..., n:, :n],
                  bounds=_bounds[..., n:, :],
              )
  
imports: from __future__ import annotations
  from copy import deepcopy
  from functools import partial
  from itertools import count
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
  from unittest.mock import patch
  import torch
  from botorch.utils.probability.linalg import PivotedCholesky
  from botorch.utils.probability.mvnxpb import MVNXPB
  from botorch.utils.testing import BotorchTestCase
  from linear_operator.utils.errors import NotPSDError
  from torch import Tensor
label_desc: Use image node_id 7 for calling the TestMVNXPB obj's underlying member method code with example usage: obj.test_augment() without return types