code: def augment_cholesky(
      Laa: Tensor,
      Kbb: Tensor,
      Kba: Optional[Tensor] = None,
      Lba: Optional[Tensor] = None,
      jitter: Optional[float] = None,
  ) -> Tensor:
      r"""Computes the Cholesky factor of a block matrix `K = [[Kaa, Kab], [Kba, Kbb]]`
      based on a precomputed Cholesky factor `Kaa = Laa Laa^T`.
  
      Args:
          Laa: Cholesky factor of K's upper left block.
          Kbb: Lower-right block of K.
          Kba: Lower-left block of K.
          Lba: Precomputed solve `Kba Laa^{-T}`.
          jitter: Optional nugget to be added to the diagonal of Kbb.
      """
      if not (Kba is None) ^ (Lba is None):
          raise ValueError(
              "One and only one of `Kba` or `Lba` must be provided."
          )
  
      if jitter is not None:
          Kbb = Kbb.clone()
          Kbb.diagonal(dim1=-2, dim2=-1).add_(jitter)
  
      if Lba is None:
          Lba = torch.linalg.solve_triangular(
              Laa.transpose(-2, -1), Kba, left=False, upper=True
          )
  
      Lbb, info = torch.linalg.cholesky_ex(
          Kbb - Lba @ Lba.transpose(-2, -1)
      )
      if info.any():
          raise NotPSDError(
              "Schur complement of `K` with respect to `Kaa` not PSD for the given "
              "Cholesky factor `Laa`"
              f"{'.' if jitter is None else f' and nugget jitter={jitter}.'}"
          )
  
      n = Lbb.shape[-1]
      return block_matrix_concat(
          blocks=([pad(Laa, (0, n))], [Lba, Lbb])
      )
  
file_path: botorch/botorch/utils/probability/linalg.py
imports: from __future__ import annotations
  from dataclasses import dataclass, InitVar
  from itertools import chain
  from typing import Any, Optional, Sequence
  import torch
  from botorch.utils.probability.utils import swap_along_dim_
  from linear_operator.utils.errors import NotPSDError
  from torch import LongTensor, Tensor
  from torch.nn.functional import pad
label_desc: Use image node_id 2 for calling a global function with example usage: augment_cholesky(Laa, Kbb, Kba, Lba, jitter) and returns: block_matrix_concat