class_docstr: An algorithm for approximating Gaussian probabilities `P(X \in bounds)`, where
  `X ~ N(0, covariance_matrix)`.
class_docstr_tok: ["An","algorithm","for","approximating","Gaussian","probabilities","`","P","(","X","\\in","bounds",")","`",",","where","`","X","~","N","(","0",",","covariance_matrix",")","`","."]
docstr: Augment an `n`-dimensional MVNXPB instance to include `m` additional random
  variables.
docstr_tok: ["Augment","an","`","n","`","-dimensional","MVNXPB","instance","to","include","`","m","`","additional","random","variables","."]
code: def augment(
      self,
      covariance_matrix: Tensor,
      bounds: Tensor,
      cross_covariance_matrix: Tensor,
      disable_pivoting: bool = False,
      jitter: Optional[float] = None,
      max_tries: Optional[int] = None,
  ) -> MVNXPB:
      r"""Augment an `n`-dimensional MVNXPB instance to include `m` additional random
      variables.
      """
      n = self.perm.shape[-1]
      m = covariance_matrix.shape[-1]
      if n != self.piv_chol.step:
          raise NotImplementedError(
              "Augmentation of incomplete solutions not implemented yet."
          )
  
      var = covariance_matrix.diagonal(dim1=-2, dim2=-1).unsqueeze(-1)
      std = var.sqrt()
      istd = var.rsqrt()
      Kmn = istd * cross_covariance_matrix
      if self.piv_chol.diag is None:
          diag = pad(
              std.squeeze(-1),
              (cross_covariance_matrix.shape[-1], 0),
              value=1,
          )
      else:
          Kmn = Kmn * (1 / self.piv_chol.diag).unsqueeze(-2)
          diag = torch.concat([self.piv_chol.diag, std.squeeze(-1)], -1)
  
      # Augment partial pivoted Cholesky factor
      Kmm = istd * covariance_matrix * istd.transpose(-1, -2)
      Lnn = self.piv_chol.tril
      try:
          L = augment_cholesky(Laa=Lnn, Kba=Kmn, Kbb=Kmm, jitter=jitter)
      except NotPSDError:
          warn(
              "Joint covariance matrix not positive definite, attempting recovery."
          )
          Knn = Lnn @ Lnn.transpose(-1, -2)
          Knm = Kmn.transpose(-1, -2)
          K = block_matrix_concat(blocks=((Knn, Knm), (Kmn, Kmm)))
          L = psd_safe_cholesky(K, jitter=jitter, max_tries=max_tries)
  
      if not disable_pivoting:
          Lmm = L[..., n:, n:].clone()
          L[..., n:, n:] = (Lmm @ Lmm.transpose(-2, -1)).tril()
  
      _bounds = istd * bounds.clip(
          *(std * lim for lim in STANDARDIZED_RANGE)
      )
      _perm = torch.arange(
          n, n + m, dtype=self.perm.dtype, device=self.perm.device
      )
      _perm = _perm.expand(covariance_matrix.shape[:-2] + (m,))
  
      piv_chol = PivotedCholesky(
          step=n + m if disable_pivoting else n,
          tril=L.contiguous(),
          perm=torch.cat(
              [self.piv_chol.perm, _perm], dim=-1
          ).contiguous(),
          diag=diag,
      )
  
      return self.build(
          step=self.step,
          perm=torch.cat([self.perm, _perm], dim=-1),
          bounds=torch.cat([self.bounds, _bounds], dim=-2),
          piv_chol=piv_chol,
          plug_ins=pad(self.plug_ins, (0, m), value=float("nan")),
          log_prob=self.log_prob,
          log_prob_extra=self.log_prob_extra,
      )
  
imports: from __future__ import annotations
  from typing import Any, Optional, TypedDict
  from warnings import warn
  import torch
  from botorch.utils.probability.bvn import bvn, bvnmom
  from botorch.utils.probability.linalg import augment_cholesky, block_matrix_concat, PivotedCholesky
  from botorch.utils.probability.utils import case_dispatcher, get_constants_like, ndtr, phi, STANDARDIZED_RANGE, swap_along_dim_
  from botorch.utils.safe_math import log, mul
  from linear_operator.utils.cholesky import psd_safe_cholesky
  from linear_operator.utils.errors import NotPSDError
  from torch import LongTensor, Tensor
  from torch.nn.functional import pad
label_desc: Use image node_id 9 for calling the MVNXPB obj's underlying member method code with example usage: obj.augment(covariance_matrix, bounds, cross_covariance_matrix, disable_pivoting, jitter, max_tries) and returns: self