class_docstr: A single task multi-fidelity GP model.
  
  A SingleTaskGP model using a DownsamplingKernel for the data fidelity
  parameter (if present) and an ExponentialDecayKernel for the iteration
  fidelity parameter (if present).
  
  This kernel is described in [Wu2019mf]_.
  
  Example:
      >>> train_X = torch.rand(20, 4)
      >>> train_Y = train_X.pow(2).sum(dim=-1, keepdim=True)
      >>> model = SingleTaskMultiFidelityGP(train_X, train_Y, data_fidelities=[3])
class_docstr_tok: ["A","single","task","multi-fidelity","GP","model",".","A","SingleTaskGP","model","using","a","DownsamplingKernel","for","the","data","fidelity","parameter","(","if","present",")","and","an","ExponentialDecayKernel","for","the","iteration","fidelity","parameter","(","if","present",")",".","This","kernel","is","described","in","[","Wu2019mf","]","_",".","Example",":",">",">",">","train_X","=","torch.rand","(","20",",","4",")",">",">",">","train_Y","=","train_X.pow","(","2",")",".sum","(","dim=-1",",","keepdim=True",")",">",">",">","model","=","SingleTaskMultiFidelityGP","(","train_X",",","train_Y",",","data_fidelities=","[","3","]",")"]
docstr: Args:
      train_X: A `batch_shape x n x (d + s)` tensor of training features,
          where `s` is the dimension of the fidelity parameters (either one
          or two).
      train_Y: A `batch_shape x n x m` tensor of training observations.
      train_Yvar: An optional `batch_shape x n x m` tensor of observed
          measurement noise.
      iteration_fidelity: The column index for the training iteration fidelity
          parameter (optional).
      data_fidelities: The column indices for the downsampling fidelity parameter.
          If a list/tuple of indices is provided, a kernel will be constructed for
          each index (optional).
      data_fidelity: The column index for the downsampling fidelity parameter
          (optional). Deprecated in favor of `data_fidelities`.
      linear_truncated: If True, use a `LinearTruncatedFidelityKernel` instead
          of the default kernel.
      nu: The smoothness parameter for the Matern kernel: either 1/2, 3/2, or
          5/2. Only used when `linear_truncated=True`.
      likelihood: A likelihood. If omitted, use a standard GaussianLikelihood
          with inferred noise level.
      outcome_transform: An outcome transform that is applied to the
              training data during instantiation and to the posterior during
              inference (that is, the `Posterior` obtained by calling
              `.posterior` on the model will be on the original scale).
      input_transform: An input transform that is applied in the model's
              forward pass.
docstr_tok: ["Args",":","train_X",":","A","`","batch_shape","x","n","x","(","d","+","s",")","`","tensor","of","training","features",",","where","`","s","`","is","the","dimension","of","the","fidelity","parameters","(","either","one","or","two",")",".","train_Y",":","A","`","batch_shape","x","n","x","m","`","tensor","of","training","observations",".","train_Yvar",":","An","optional","`","batch_shape","x","n","x","m","`","tensor","of","observed","measurement","noise",".","iteration_fidelity",":","The","column","index","for","the","training","iteration","fidelity","parameter","(","optional",")",".","data_fidelities",":","The","column","indices","for","the","downsampling","fidelity","parameter",".","If","a","list\/tuple","of","indices","is","provided",",","a","kernel","will","be","constructed","for","each","index","(","optional",")",".","data_fidelity",":","The","column","index","for","the","downsampling","fidelity","parameter","(","optional",")",".","Deprecated","in","favor","of","`","data_fidelities","`",".","linear_truncated",":","If","True",",","use","a","`","LinearTruncatedFidelityKernel","`","instead","of","the","default","kernel",".","nu",":","The","smoothness","parameter","for","the","Matern","kernel",":","either","1\/2",",","3\/2",",","or","5\/2",".","Only","used","when","`","linear_truncated=True","`",".","likelihood",":","A","likelihood",".","If","omitted",",","use","a","standard","GaussianLikelihood","with","inferred","noise","level",".","outcome_transform",":","An","outcome","transform","that","is","applied","to","the","training","data","during","instantiation","and","to","the","posterior","during","inference","(","that","is",",","the","`","Posterior","`","obtained","by","calling","`",".posterior","`","on","the","model","will","be","on","the","original","scale",")",".","input_transform",":","An","input","transform","that","is","applied","in","the","model's","forward","pass","."]
code: def __init__(
      self,
      train_X: Tensor,
      train_Y: Tensor,
      train_Yvar: Optional[Tensor] = None,
      iteration_fidelity: Optional[int] = None,
      data_fidelities: Optional[Union[List[int], Tuple[int]]] = None,
      data_fidelity: Optional[int] = None,
      linear_truncated: bool = True,
      nu: float = 2.5,
      likelihood: Optional[Likelihood] = None,
      outcome_transform: Optional[OutcomeTransform] = None,
      input_transform: Optional[InputTransform] = None,
  ) -> None:
      r"""
      Args:
          train_X: A `batch_shape x n x (d + s)` tensor of training features,
              where `s` is the dimension of the fidelity parameters (either one
              or two).
          train_Y: A `batch_shape x n x m` tensor of training observations.
          train_Yvar: An optional `batch_shape x n x m` tensor of observed
              measurement noise.
          iteration_fidelity: The column index for the training iteration fidelity
              parameter (optional).
          data_fidelities: The column indices for the downsampling fidelity parameter.
              If a list/tuple of indices is provided, a kernel will be constructed for
              each index (optional).
          data_fidelity: The column index for the downsampling fidelity parameter
              (optional). Deprecated in favor of `data_fidelities`.
          linear_truncated: If True, use a `LinearTruncatedFidelityKernel` instead
              of the default kernel.
          nu: The smoothness parameter for the Matern kernel: either 1/2, 3/2, or
              5/2. Only used when `linear_truncated=True`.
          likelihood: A likelihood. If omitted, use a standard GaussianLikelihood
              with inferred noise level.
          outcome_transform: An outcome transform that is applied to the
                  training data during instantiation and to the posterior during
                  inference (that is, the `Posterior` obtained by calling
                  `.posterior` on the model will be on the original scale).
          input_transform: An input transform that is applied in the model's
                  forward pass.
      """
      if data_fidelity is not None:
          warnings.warn(
              "The `data_fidelity` argument is deprecated and will be removed in "
              "a future release. Please use `data_fidelities` instead.",
              DeprecationWarning,
          )
          if data_fidelities is not None:
              raise ValueError(
                  "Cannot specify both `data_fidelity` and `data_fidelities`."
              )
          data_fidelities = [data_fidelity]
  
      self._init_args = {
          "iteration_fidelity": iteration_fidelity,
          "data_fidelities": data_fidelities,
          "linear_truncated": linear_truncated,
          "nu": nu,
          "outcome_transform": outcome_transform,
      }
      if iteration_fidelity is None and data_fidelities is None:
          raise UnsupportedError(
              "SingleTaskMultiFidelityGP requires at least one fidelity parameter."
          )
      with torch.no_grad():
          transformed_X = self.transform_inputs(
              X=train_X, input_transform=input_transform
          )
  
      self._set_dimensions(train_X=transformed_X, train_Y=train_Y)
      (
          covar_module,
          subset_batch_dict,
      ) = _setup_multifidelity_covar_module(
          dim=transformed_X.size(-1),
          aug_batch_shape=self._aug_batch_shape,
          iteration_fidelity=iteration_fidelity,
          data_fidelities=data_fidelities,
          linear_truncated=linear_truncated,
          nu=nu,
      )
      super().__init__(
          train_X=train_X,
          train_Y=train_Y,
          train_Yvar=train_Yvar,
          likelihood=likelihood,
          covar_module=covar_module,
          outcome_transform=outcome_transform,
          input_transform=input_transform,
      )
      self._subset_batch_dict = {
          "mean_module.raw_constant": -1,
          "covar_module.raw_outputscale": -1,
          **subset_batch_dict,
      }
      if train_Yvar is None:
          self._subset_batch_dict[
              "likelihood.noise_covar.raw_noise"
          ] = -2
      self.to(train_X)
  
imports: from __future__ import annotations
  import warnings
  from typing import Any, Dict, List, Optional, Tuple, Union
  import torch
  from botorch.exceptions.errors import UnsupportedError
  from botorch.models.gp_regression import SingleTaskGP
  from botorch.models.kernels.downsampling import DownsamplingKernel
  from botorch.models.kernels.exponential_decay import ExponentialDecayKernel
  from botorch.models.kernels.linear_truncated_fidelity import LinearTruncatedFidelityKernel
  from botorch.models.transforms.input import InputTransform
  from botorch.models.transforms.outcome import OutcomeTransform
  from botorch.utils.datasets import SupervisedDataset
  from gpytorch.kernels.kernel import ProductKernel
  from gpytorch.kernels.rbf_kernel import RBFKernel
  from gpytorch.kernels.scale_kernel import ScaleKernel
  from gpytorch.likelihoods.likelihood import Likelihood
  from gpytorch.priors.torch_priors import GammaPrior
  from torch import Tensor
label_desc: Use image node_id 1 to create a new SingleTaskMultiFidelityGP object from inherited base classes: SingleTaskGP with example: obj = SingleTaskMultiFidelityGP(train_X, train_Y, train_Yvar, iteration_fidelity, data_fidelities, data_fidelity, linear_truncated, nu, likelihood, outcome_transform, input_transform)