code: def _setup_multifidelity_covar_module(
      dim: int,
      aug_batch_shape: torch.Size,
      iteration_fidelity: Optional[int],
      data_fidelities: Optional[List[int]],
      linear_truncated: bool,
      nu: float,
  ) -> Tuple[ScaleKernel, Dict]:
      """Helper function to get the covariance module and associated subset_batch_dict
      for the multifidelity setting.
  
      Args:
          dim: The dimensionality of the training data.
          aug_batch_shape: The output-augmented batch shape as defined in
              `BatchedMultiOutputGPyTorchModel`.
          iteration_fidelity: The column index for the training iteration fidelity
              parameter (optional).
          data_fidelities: The column indices for the downsampling fidelity parameters
              (optional).
          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`.
  
      Returns:
          The covariance module and subset_batch_dict.
      """
  
      if iteration_fidelity is not None and iteration_fidelity < 0:
          iteration_fidelity = dim + iteration_fidelity
      if data_fidelities is not None:
          for i in range(len(data_fidelities)):
              if data_fidelities[i] < 0:
                  data_fidelities[i] = dim + data_fidelities[i]
  
      kernels = []
  
      if linear_truncated:
          leading_dims = (
              [iteration_fidelity]
              if iteration_fidelity is not None
              else []
          )
          trailing_dims = (
              [[i] for i in data_fidelities]
              if data_fidelities is not None
              else [[]]
          )
          for tdims in trailing_dims:
              kernels.append(
                  LinearTruncatedFidelityKernel(
                      fidelity_dims=leading_dims + tdims,
                      dimension=dim,
                      nu=nu,
                      batch_shape=aug_batch_shape,
                      power_prior=GammaPrior(3.0, 3.0),
                  )
              )
      else:
          non_active_dims = set(data_fidelities or [])
          if iteration_fidelity is not None:
              non_active_dims.add(iteration_fidelity)
          active_dimsX = sorted(set(range(dim)) - non_active_dims)
          kernels.append(
              RBFKernel(
                  ard_num_dims=len(active_dimsX),
                  batch_shape=aug_batch_shape,
                  lengthscale_prior=GammaPrior(3.0, 6.0),
                  active_dims=active_dimsX,
              )
          )
          if iteration_fidelity is not None:
              kernels.append(
                  ExponentialDecayKernel(
                      batch_shape=aug_batch_shape,
                      lengthscale_prior=GammaPrior(3.0, 6.0),
                      offset_prior=GammaPrior(3.0, 6.0),
                      power_prior=GammaPrior(3.0, 6.0),
                      active_dims=[iteration_fidelity],
                  )
              )
          if data_fidelities is not None:
              for data_fidelity in data_fidelities:
                  kernels.append(
                      DownsamplingKernel(
                          batch_shape=aug_batch_shape,
                          offset_prior=GammaPrior(3.0, 6.0),
                          power_prior=GammaPrior(3.0, 6.0),
                          active_dims=[data_fidelity],
                      )
                  )
  
      kernel = ProductKernel(*kernels)
  
      covar_module = ScaleKernel(
          kernel,
          batch_shape=aug_batch_shape,
          outputscale_prior=GammaPrior(2.0, 0.15),
      )
  
      key_prefix = "covar_module.base_kernel.kernels"
      if linear_truncated:
          subset_batch_dict = {}
          for i in range(len(kernels)):
              subset_batch_dict.update(
                  {
                      f"{key_prefix}.{i}.raw_power": -2,
                      f"{key_prefix}.{i}.covar_module_unbiased.raw_lengthscale": -3,
                      f"{key_prefix}.{i}.covar_module_biased.raw_lengthscale": -3,
                  }
              )
      else:
          subset_batch_dict = {
              f"{key_prefix}.0.raw_lengthscale": -3,
          }
  
          if iteration_fidelity is not None:
              subset_batch_dict.update(
                  {
                      f"{key_prefix}.1.raw_power": -2,
                      f"{key_prefix}.1.raw_offset": -2,
                      f"{key_prefix}.1.raw_lengthscale": -3,
                  }
              )
          if data_fidelities is not None:
              start_idx = 2 if iteration_fidelity is not None else 1
              for i in range(
                  start_idx, len(data_fidelities) + start_idx
              ):
                  subset_batch_dict.update(
                      {
                          f"{key_prefix}.{i}.raw_power": -2,
                          f"{key_prefix}.{i}.raw_offset": -2,
                      }
                  )
  
      return covar_module, subset_batch_dict
  
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 for calling a global function with example usage: _setup_multifidelity_covar_module(dim, aug_batch_shape, iteration_fidelity, data_fidelities, linear_truncated, nu) and returns: covar_module,  subset_batch_dict