class_docstr: A fully Bayesian single-task GP model with the SAAS prior.
  
  This model assumes that the inputs have been normalized to [0, 1]^d and that
  the output has been standardized to have zero mean and unit variance. You can
  either normalize and standardize the data before constructing the model or use
  an `input_transform` and `outcome_transform`. The SAAS model [Eriksson2021saasbo]_
  with a Matern-5/2 kernel is used by default.
  
  You are expected to use `fit_fully_bayesian_model_nuts` to fit this model as it
  isn't compatible with `fit_gpytorch_model`.
  
  Example:
      >>> saas_gp = SaasFullyBayesianSingleTaskGP(train_X, train_Y)
      >>> fit_fully_bayesian_model_nuts(saas_gp)
      >>> posterior = saas_gp.posterior(test_X)
class_docstr_tok: ["A","fully","Bayesian","single-task","GP","model","with","the","SAAS","prior",".","This","model","assumes","that","the","inputs","have","been","normalized","to","[","0",",","1","]","^d","and","that","the","output","has","been","standardized","to","have","zero","mean","and","unit","variance",".","You","can","either","normalize","and","standardize","the","data","before","constructing","the","model","or","use","an","`","input_transform","`","and","`","outcome_transform","`",".","The","SAAS","model","[","Eriksson2021saasbo","]","_","with","a","Matern-5\/2","kernel","is","used","by","default",".","You","are","expected","to","use","`","fit_fully_bayesian_model_nuts","`","to","fit","this","model","as","it","is","n't","compatible","with","`","fit_gpytorch_model","`",".","Example",":",">",">",">","saas_gp","=","SaasFullyBayesianSingleTaskGP","(","train_X",",","train_Y",")",">",">",">","fit_fully_bayesian_model_nuts","(","saas_gp",")",">",">",">","posterior","=","saas_gp.posterior","(","test_X",")"]
docstr: Custom logic for loading the state dict.
  
  The standard approach of calling `load_state_dict` currently doesn't play well
  with the `SaasFullyBayesianSingleTaskGP` since the `mean_module`, `covar_module`
  and `likelihood` aren't initialized until the model has been fitted. The reason
  for this is that we don't know the number of MCMC samples until NUTS is called.
  Given the state dict, we can initialize a new model with some dummy samples and
  then load the state dict into this model. This currently only works for a
  `SaasPyroModel` and supporting more Pyro models likely requires moving the model
  construction logic into the Pyro model itself.
docstr_tok: ["Custom","logic","for","loading","the","state","dict",".","The","standard","approach","of","calling","`","load_state_dict","`","currently","does","n't","play","well","with","the","`","SaasFullyBayesianSingleTaskGP","`","since","the","`","mean_module","`",",","`","covar_module","`","and","`","likelihood","`","are","n't","initialized","until","the","model","has","been","fitted",".","The","reason","for","this","is","that","we","do","n't","know","the","number","of","MCMC","samples","until","NUTS","is","called",".","Given","the","state","dict",",","we","can","initialize","a","new","model","with","some","dummy","samples","and","then","load","the","state","dict","into","this","model",".","This","currently","only","works","for","a","`","SaasPyroModel","`","and","supporting","more","Pyro","models","likely","requires","moving","the","model","construction","logic","into","the","Pyro","model","itself","."]
code: def load_state_dict(
      self, state_dict: Mapping[str, Any], strict: bool = True
  ):
      r"""Custom logic for loading the state dict.
  
      The standard approach of calling `load_state_dict` currently doesn't play well
      with the `SaasFullyBayesianSingleTaskGP` since the `mean_module`, `covar_module`
      and `likelihood` aren't initialized until the model has been fitted. The reason
      for this is that we don't know the number of MCMC samples until NUTS is called.
      Given the state dict, we can initialize a new model with some dummy samples and
      then load the state dict into this model. This currently only works for a
      `SaasPyroModel` and supporting more Pyro models likely requires moving the model
      construction logic into the Pyro model itself.
      """
  
      if not isinstance(self.pyro_model, SaasPyroModel):
          raise NotImplementedError(
              "load_state_dict only works for SaasPyroModel"
          )
      raw_mean = state_dict["mean_module.raw_constant"]
      num_mcmc_samples = len(raw_mean)
      dim = self.pyro_model.train_X.shape[-1]
      tkwargs = {"device": raw_mean.device, "dtype": raw_mean.dtype}
      # Load some dummy samples
      mcmc_samples = {
          "mean": torch.ones(num_mcmc_samples, **tkwargs),
          "lengthscale": torch.ones(num_mcmc_samples, dim, **tkwargs),
          "outputscale": torch.ones(num_mcmc_samples, **tkwargs),
      }
      if self.pyro_model.train_Yvar is None:
          mcmc_samples["noise"] = torch.ones(
              num_mcmc_samples, **tkwargs
          )
      (
          self.mean_module,
          self.covar_module,
          self.likelihood,
      ) = self.pyro_model.load_mcmc_samples(mcmc_samples=mcmc_samples)
      # Load the actual samples from the state dict
      super().load_state_dict(state_dict=state_dict, strict=strict)
  
imports: import math
  from abc import abstractmethod
  from typing import Any, Dict, List, Mapping, Optional, Tuple
  import pyro
  import torch
  from botorch.acquisition.objective import PosteriorTransform
  from botorch.models.gpytorch import BatchedMultiOutputGPyTorchModel
  from botorch.models.transforms.input import InputTransform
  from botorch.models.transforms.outcome import OutcomeTransform
  from botorch.models.utils import validate_input_scaling
  from botorch.models.utils.gpytorch_modules import MIN_INFERRED_NOISE_LEVEL
  from botorch.posteriors.fully_bayesian import GaussianMixturePosterior, MCMC_DIM
  from gpytorch.constraints import GreaterThan
  from gpytorch.distributions.multivariate_normal import MultivariateNormal
  from gpytorch.kernels import MaternKernel, ScaleKernel
  from gpytorch.kernels.kernel import dist, Kernel
  from gpytorch.likelihoods.gaussian_likelihood import FixedNoiseGaussianLikelihood, GaussianLikelihood
  from gpytorch.likelihoods.likelihood import Likelihood
  from gpytorch.means.constant_mean import ConstantMean
  from gpytorch.means.mean import Mean
  from gpytorch.models.exact_gp import ExactGP
  from pyro.ops.integrator import register_exception_handler
  from torch import Tensor
label_desc: Use image node_id 9 for calling the SaasFullyBayesianSingleTaskGP obj's underlying member method code with example usage: obj.load_state_dict(state_dict, strict) without return types