class_docstr: Multi-task GP with Kronecker structure, using an ICM kernel.
  
  This model assumes the "block design" case, i.e., it requires that all tasks
  are observed at all data points.
  
  For posterior sampling, this model uses Matheron's rule [Doucet2010sampl] to compute
  the posterior over all tasks as in [Maddox2021bohdo] by exploiting Kronecker
  structure.
  
  When a multi-fidelity model has Kronecker structure, this means there is one
  covariance kernel over the fidelity features (call it `K_f`) and another over
  the rest of the input parameters (call it `K_i`), and the resulting covariance
  across inputs and fidelities is given by the Kronecker product of the two
  covariance matrices. This is equivalent to saying the covariance between
  two input and feature pairs is given by
  
  K((parameter_1, fidelity_1), (parameter_2, fidelity_2))
      = K_f(fidelity_1, fidelity_2) * K_i(parameter_1, parameter_2).
  
  Then the covariance matrix of `n_i` parameters and `n_f` fidelities can be
  codified as a Kronecker product of an `n_i x n_i` matrix and an
  `n_f x n_f` matrix, which is far more parsimonious than specifying the
  whole `(n_i * n_f) x (n_i * n_f)` covariance matrix.
  
  Example:
      >>> train_X = torch.rand(10, 2)
      >>> train_Y = torch.cat([f_1(X), f_2(X)], dim=-1)
      >>> model = KroneckerMultiTaskGP(train_X, train_Y)
class_docstr_tok: ["Multi-task","GP","with","Kronecker","structure",",","using","an","ICM","kernel",".","This","model","assumes","the","``","block","design","''","case",",","i.e.",",","it","requires","that","all","tasks","are","observed","at","all","data","points",".","For","posterior","sampling",",","this","model","uses","Matheron","'s","rule","[","Doucet2010sampl","]","to","compute","the","posterior","over","all","tasks","as","in","[","Maddox2021bohdo","]","by","exploiting","Kronecker","structure",".","When","a","multi-fidelity","model","has","Kronecker","structure",",","this","means","there","is","one","covariance","kernel","over","the","fidelity","features","(","call","it","`","K_f","`",")","and","another","over","the","rest","of","the","input","parameters","(","call","it","`","K_i","`",")",",","and","the","resulting","covariance","across","inputs","and","fidelities","is","given","by","the","Kronecker","product","of","the","two","covariance","matrices",".","This","is","equivalent","to","saying","the","covariance","between","two","input","and","feature","pairs","is","given","by","K","(","(","parameter_1",",","fidelity_1",")",",","(","parameter_2",",","fidelity_2",")",")","=","K_f","(","fidelity_1",",","fidelity_2",")","*","K_i","(","parameter_1",",","parameter_2",")",".","Then","the","covariance","matrix","of","`","n_i","`","parameters","and","`","n_f","`","fidelities","can","be","codified","as","a","Kronecker","product","of","an","`","n_i","x","n_i","`","matrix","and","an","`","n_f","x","n_f","`","matrix",",","which","is","far","more","parsimonious","than","specifying","the","whole","`","(","n_i","*","n_f",")","x","(","n_i","*","n_f",")","`","covariance","matrix",".","Example",":",">",">",">","train_X","=","torch.rand","(","10",",","2",")",">",">",">","train_Y","=","torch.cat","(","[","f_1","(","X",")",",","f_2","(","X",")","]",",","dim=-1",")",">",">",">","model","=","KroneckerMultiTaskGP","(","train_X",",","train_Y",")"]
code: def posterior(
      self,
      X: Tensor,
      output_indices: Optional[List[int]] = None,
      observation_noise: Union[bool, Tensor] = False,
      posterior_transform: Optional[PosteriorTransform] = None,
      **kwargs: Any,
  ) -> MultitaskGPPosterior:
      self.eval()
  
      if posterior_transform is not None:
          # this could be very costly, disallow for now
          raise NotImplementedError(
              "Posterior transforms currently not supported for "
              f"{self.__class__.__name__}"
          )
  
      X = self.transform_inputs(X)
      train_x = self.transform_inputs(self.train_inputs[0])
  
      # construct Ktt
      task_covar = self._task_covar_matrix
      task_rootlt = self._task_covar_matrix.root_decomposition(
          method="diagonalization"
      )
      task_root = task_rootlt.root
      if task_covar.batch_shape != X.shape[:-2]:
          task_covar = BatchRepeatLinearOperator(
              task_covar, batch_repeat=X.shape[:-2]
          )
          task_root = BatchRepeatLinearOperator(
              to_linear_operator(task_root), batch_repeat=X.shape[:-2]
          )
  
      task_covar_rootlt = RootLinearOperator(task_root)
  
      # construct RR' \approx Kxx
      data_data_covar = self.train_full_covar.linear_ops[0]
      # populate the diagonalziation caches for the root and inverse root
      # decomposition
      (
          data_data_evals,
          data_data_evecs,
      ) = data_data_covar.diagonalization()
  
      # pad the eigenvalue and eigenvectors with zeros if we are using lanczos
      if data_data_evecs.shape[-1] < data_data_evecs.shape[-2]:
          cols_to_add = (
              data_data_evecs.shape[-2] - data_data_evecs.shape[-1]
          )
          zero_evecs = torch.zeros(
              *data_data_evecs.shape[:-1],
              cols_to_add,
              dtype=data_data_evals.dtype,
              device=data_data_evals.device,
          )
          zero_evals = torch.zeros(
              *data_data_evecs.shape[:-2],
              cols_to_add,
              dtype=data_data_evals.dtype,
              device=data_data_evals.device,
          )
          data_data_evecs = CatLinearOperator(
              data_data_evecs,
              to_linear_operator(zero_evecs),
              dim=-1,
              output_device=data_data_evals.device,
          )
          data_data_evals = torch.cat(
              (data_data_evals, zero_evals), dim=-1
          )
  
      # construct K_{xt, x}
      test_data_covar = self.covar_module.data_covar_module(X, train_x)
      # construct K_{xt, xt}
      test_test_covar = self.covar_module.data_covar_module(X)
  
      # now update root so that \tilde{R}\tilde{R}' \approx K_{(x,xt), (x,xt)}
      # cloning preserves the gradient history
      updated_linear_op = data_data_covar.cat_rows(
          cross_mat=test_data_covar.clone(),
          new_mat=test_test_covar,
          method="diagonalization",
      )
      updated_root = updated_linear_op.root_decomposition().root
      # occasionally, there's device errors so enforce this comes out right
      updated_root = updated_root.to(data_data_covar.device)
  
      # build a root decomposition of the joint train/test covariance matrix
      # construct (\tilde{R} \otimes M)(\tilde{R} \otimes M)' \approx
      # (K_{(x,xt), (x,xt)} \otimes Ktt)
      joint_covar = RootLinearOperator(
          KroneckerProductLinearOperator(
              updated_root, task_covar_rootlt.root.detach()
          )
      )
  
      # construct K_{xt, x} \otimes Ktt
      test_obs_kernel = KroneckerProductLinearOperator(
          test_data_covar, task_covar
      )
  
      # collect y - \mu(x) and \mu(X)
      train_diff = self.train_targets - self.mean_module(train_x)
      if detach_test_caches.on():
          train_diff = train_diff.detach()
      test_mean = self.mean_module(X)
  
      train_noise = self.likelihood._shaped_noise_covar(train_x.shape)
      diagonal_noise = isinstance(train_noise, DiagLinearOperator)
      if detach_test_caches.on():
          train_noise = train_noise.detach()
      test_noise = (
          self.likelihood._shaped_noise_covar(X.shape)
          if observation_noise
          else None
      )
  
      # predictive mean and variance for the mvn
      # first the predictive mean
      pred_mean = (
          test_obs_kernel.matmul(self.predictive_mean_cache).reshape_as(
              test_mean
          )
          + test_mean
      )
      # next the predictive variance, assume diagonal noise
      test_var_term = KroneckerProductLinearOperator(
          test_test_covar, task_covar
      ).diagonal()
  
      if diagonal_noise:
          (
              task_evals,
              task_evecs,
          ) = self._task_covar_matrix.diagonalization()
          # TODO: make this be the default KPMatmulLT diagonal method in gpytorch
          full_data_inv_evals = (
              KroneckerProductDiagLinearOperator(
                  DiagLinearOperator(data_data_evals),
                  DiagLinearOperator(task_evals),
              )
              + train_noise
          ).inverse()
          test_train_hadamard = KroneckerProductLinearOperator(
              test_data_covar.matmul(data_data_evecs).to_dense() ** 2,
              task_covar.matmul(task_evecs).to_dense() ** 2,
          )
          data_var_term = test_train_hadamard.matmul(
              full_data_inv_evals
          ).sum(dim=-1)
      else:
          # if non-diagonal noise (but still kronecker structured), we have to pull
          # across the noise because the inverse is not closed form
          # should be a kronecker lt, R = \Sigma_X^{-1/2} \kron \Sigma_T^{-1/2}
          # TODO: enforce the diagonalization to return a KPLT for all shapes in
          # gpytorch or dense linear algebra for small shapes
          data_noise, task_noise = train_noise.linear_ops
          data_noise_root = data_noise.root_inv_decomposition(
              method="diagonalization"
          )
          task_noise_root = task_noise.root_inv_decomposition(
              method="diagonalization"
          )
  
          # ultimately we need to compute the diagonal of
          # (K_{x* X} \kron K_T)(K_{XX} \kron K_T + \Sigma_X \kron \Sigma_T)^{-1}
          #                           (K_{x* X} \kron K_T)^T
          # = (K_{x* X} \Sigma_X^{-1/2} Q_R)(\Lambda_R + I)^{-1}
          #                       (K_{x* X} \Sigma_X^{-1/2} Q_R)^T
          # where R = (\Sigma_X^{-1/2T}K_{XX}\Sigma_X^{-1/2} \kron
          #                   \Sigma_T^{-1/2T}K_{T}\Sigma_T^{-1/2})
          # first we construct the components of R's eigen-decomposition
          # TODO: make this be the default KPMatmulLT diagonal method in gpytorch
          whitened_data_covar = (
              data_noise_root.transpose(-1, -2)
              .matmul(data_data_covar)
              .matmul(data_noise_root)
          )
          (
              w_data_evals,
              w_data_evecs,
          ) = whitened_data_covar.diagonalization()
          whitened_task_covar = (
              task_noise_root.transpose(-1, -2)
              .matmul(self._task_covar_matrix)
              .matmul(task_noise_root)
          )
          (
              w_task_evals,
              w_task_evecs,
          ) = whitened_task_covar.diagonalization()
  
          # we add one to the eigenvalues as above (not just for stability)
          full_data_inv_evals = (
              KroneckerProductDiagLinearOperator(
                  DiagLinearOperator(w_data_evals),
                  DiagLinearOperator(w_task_evals),
              )
              .add_jitter(1.0)
              .inverse()
          )
  
          test_data_comp = (
              test_data_covar.matmul(data_noise_root)
              .matmul(w_data_evecs)
              .to_dense()
              ** 2
          )
          task_comp = (
              task_covar.matmul(task_noise_root)
              .matmul(w_task_evecs)
              .to_dense()
              ** 2
          )
  
          test_train_hadamard = KroneckerProductLinearOperator(
              test_data_comp, task_comp
          )
          data_var_term = test_train_hadamard.matmul(
              full_data_inv_evals
          ).sum(dim=-1)
  
      pred_variance = test_var_term - data_var_term
      specialized_mvn = MultitaskMultivariateNormal(
          pred_mean, DiagLinearOperator(pred_variance)
      )
      if observation_noise:
          specialized_mvn = self.likelihood(specialized_mvn)
  
      posterior = MultitaskGPPosterior(
          distribution=specialized_mvn,
          joint_covariance_matrix=joint_covar,
          test_train_covar=test_obs_kernel,
          train_diff=train_diff,
          test_mean=test_mean,
          train_train_covar=self.train_full_covar,
          train_noise=train_noise,
          test_noise=test_noise,
      )
  
      if hasattr(self, "outcome_transform"):
          posterior = self.outcome_transform.untransform_posterior(
              posterior
          )
      return posterior
  
imports: from __future__ import annotations
  import math
  import warnings
  from typing import Any, Dict, List, Optional, Tuple, Union
  import torch
  from botorch.acquisition.objective import PosteriorTransform
  from botorch.models.gpytorch import GPyTorchModel, MultiTaskGPyTorchModel
  from botorch.models.model import FantasizeMixin
  from botorch.models.transforms.input import InputTransform
  from botorch.models.transforms.outcome import OutcomeTransform
  from botorch.models.utils.gpytorch_modules import get_matern_kernel_with_gamma_prior, MIN_INFERRED_NOISE_LEVEL
  from botorch.posteriors.multitask import MultitaskGPPosterior
  from botorch.utils.datasets import MultiTaskDataset, SupervisedDataset
  from gpytorch.constraints import GreaterThan
  from gpytorch.distributions.multitask_multivariate_normal import MultitaskMultivariateNormal
  from gpytorch.distributions.multivariate_normal import MultivariateNormal
  from gpytorch.kernels.index_kernel import IndexKernel
  from gpytorch.kernels.matern_kernel import MaternKernel
  from gpytorch.kernels.multitask_kernel import MultitaskKernel
  from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood, GaussianLikelihood
  from gpytorch.likelihoods.likelihood import Likelihood
  from gpytorch.likelihoods.multitask_gaussian_likelihood import MultitaskGaussianLikelihood
  from gpytorch.means import MultitaskMean
  from gpytorch.means.constant_mean import ConstantMean
  from gpytorch.models.exact_gp import ExactGP
  from gpytorch.module import Module
  from gpytorch.priors.lkj_prior import LKJCovariancePrior
  from gpytorch.priors.prior import Prior
  from gpytorch.priors.smoothed_box_prior import SmoothedBoxPrior
  from gpytorch.priors.torch_priors import GammaPrior
  from gpytorch.settings import detach_test_caches
  from gpytorch.utils.errors import CachingError
  from gpytorch.utils.memoize import cached, pop_from_cache
  from linear_operator.operators import BatchRepeatLinearOperator, CatLinearOperator, DiagLinearOperator, KroneckerProductDiagLinearOperator, KroneckerProductLinearOperator, RootLinearOperator, to_linear_operator
  from torch import Tensor
label_desc: Use image node_id 6 for calling the KroneckerMultiTaskGP obj's underlying member method code with example usage: obj.posterior(X, output_indices, observation_noise, posterior_transform) and returns: posterior