class_docstr: The GP for ALEBO.
  
  Uses the Mahalanobis kernel defined in ALEBOKernel, along with a
  ScaleKernel to add a kernel variance and a fitted constant mean.
  
  In non-batch mode, there is a single kernel that produces MVN predictions
  as usual for a GP.
  With b batches, each batch has its own set of kernel hyperparameters and
  each batch represents a sample from the hyperparameter posterior
  distribution. When making a prediction (with `__call__`), these samples are
  integrated over using moment matching. So, the predictions are an MVN as
  usual with the same shape as in non-batch mode.
  
  Args:
      B: (d x D) Projection matrix.
      train_X: (n x d) X training data.
      train_Y: (n x 1) Y training data.
      train_Yvar: (n x 1) Noise variances of each training Y.
class_docstr_tok: ["The","GP","for","ALEBO",".","Uses","the","Mahalanobis","kernel","defined","in","ALEBOKernel",",","along","with","a","ScaleKernel","to","add","a","kernel","variance","and","a","fitted","constant","mean",".","In","non-batch","mode",",","there","is","a","single","kernel","that","produces","MVN","predictions","as","usual","for","a","GP",".","With","b","batches",",","each","batch","has","its","own","set","of","kernel","hyperparameters","and","each","batch","represents","a","sample","from","the","hyperparameter","posterior","distribution",".","When","making","a","prediction","(","with","`","__call__","`",")",",","these","samples","are","integrated","over","using","moment","matching",".","So",",","the","predictions","are","an","MVN","as","usual","with","the","same","shape","as","in","non-batch","mode",".","Args",":","B",":","(","d","x","D",")","Projection","matrix",".","train_X",":","(","n","x","d",")","X","training","data",".","train_Y",":","(","n","x","1",")","Y","training","data",".","train_Yvar",":","(","n","x","1",")","Noise","variances","of","each","training","Y","."]
docstr: If model is non-batch, then just make a prediction. If model has
  multiple batches, then these are samples from the kernel hyperparameter
  posterior and we integrate over them with moment matching.
  
  The shape of the MVN that this outputs will be the same regardless of
  whether the model is batched or not.
  
  Args:
      x: Point to be predicted.
  
  Returns: MultivariateNormal distribution of prediction.
docstr_tok: ["If","model","is","non-batch",",","then","just","make","a","prediction",".","If","model","has","multiple","batches",",","then","these","are","samples","from","the","kernel","hyperparameter","posterior","and","we","integrate","over","them","with","moment","matching",".","The","shape","of","the","MVN","that","this","outputs","will","be","the","same","regardless","of","whether","the","model","is","batched","or","not",".","Args",":","x",":","Point","to","be","predicted",".","Returns",":","MultivariateNormal","distribution","of","prediction","."]
code: def __call__(self, x: Tensor) -> MultivariateNormal:
      """
      If model is non-batch, then just make a prediction. If model has
      multiple batches, then these are samples from the kernel hyperparameter
      posterior and we integrate over them with moment matching.
  
      The shape of the MVN that this outputs will be the same regardless of
      whether the model is batched or not.
  
      Args:
          x: Point to be predicted.
  
      Returns: MultivariateNormal distribution of prediction.
      """
      if len(self._aug_batch_shape) == 0:
          return super().__call__(x)
      # Else, approximately integrate over batches with moment matching.
      # Take X as (b) x q x d, and expand to (b) x ns x q x d
      if x.ndim > 3:
          raise ValueError("Don't know how to predict this shape")
      x = x.unsqueeze(-3).expand(
          x.shape[:-2]
          + torch.Size([self._aug_batch_shape[0]])  # pyre-ignore
          # pyre-fixme[58]: `+` is not supported for operand types `Tuple[int,
          #  ...]` and `Size`.
          + x.shape[-2:]
      )
      mvn_b = super().__call__(x)
      mu = mvn_b.mean.mean(dim=-2)
      C = (
          mvn_b.covariance_matrix.mean(dim=-3)
          + torch.matmul(mvn_b.mean.transpose(-2, -1), mvn_b.mean)
          / mvn_b.mean.shape[-2]
          - torch.matmul(mu.unsqueeze(-1), mu.unsqueeze(-2))
      )  # Law of Total Covariance
      mvn = MultivariateNormal(mu, C)
      return mvn
  
imports: from __future__ import annotations
  import dataclasses
  import re
  from collections import OrderedDict
  from logging import Logger
  from math import inf
  from numbers import Number
  from typing import Any, Callable, Dict, List, MutableMapping, Optional, Tuple, TypeVar, Union
  from warnings import warn
  import gpytorch
  import numpy
  import torch
  from ax.core.search_space import SearchSpaceDigest
  from ax.core.types import TCandidateMetadata
  from ax.models.random.alebo_initializer import ALEBOInitializer
  from ax.models.torch.botorch import BotorchModel
  from ax.models.torch.botorch_defaults import get_qLogNEI
  from ax.models.torch.utils import _datasets_to_legacy_inputs
  from ax.models.torch_base import TorchGenResults, TorchModel, TorchOptConfig
  from ax.utils.common.docutils import copy_doc
  from ax.utils.common.logger import get_logger
  from ax.utils.common.typeutils import checked_cast
  from botorch.acquisition.acquisition import AcquisitionFunction
  from botorch.acquisition.analytic import ExpectedImprovement
  from botorch.acquisition.objective import PosteriorTransform
  from botorch.models.gp_regression import SingleTaskGP
  from botorch.models.gpytorch import GPyTorchModel
  from botorch.models.model_list_gp_regression import ModelListGP
  from botorch.optim.fit import fit_gpytorch_mll_scipy
  from botorch.optim.initializers import initialize_q_batch_nonneg
  from botorch.optim.optimize import optimize_acqf
  from botorch.optim.utils import _handle_numerical_errors, get_parameters_and_bounds, TorchAttr
  from botorch.posteriors.gpytorch import GPyTorchPosterior
  from botorch.utils.datasets import SupervisedDataset
  from gpytorch.distributions.multivariate_normal import MultivariateNormal
  from gpytorch.kernels.kernel import Kernel
  from gpytorch.kernels.rbf_kernel import postprocess_rbf
  from gpytorch.kernels.scale_kernel import ScaleKernel
  from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
  from scipy.optimize import approx_fprime
  from torch import Tensor
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 2 for calling the ALEBOGP obj's underlying member method code with example usage: obj.__call__(x) and returns: mvn, super