class_docstr: A GP using a Latent Context Embedding Additive (LCE-A) Kernel.
  
  Note that the model does not support batch training. Input training
  data sets should have dim = 2.
docstr: Args:
      train_X: (n x d) X training data.
      train_Y: (n x 1) Y training data.
      train_Yvar: (n x 1) Noise variance of Y. If None,
          we use an inferred noise likelihood.
      decomposition: Keys are context names. Values are the indexes of
          parameters belong to the context.
      train_embedding: Whether to train the embedding layer or not. If False,
          the model will use pre-trained embeddings in embs_feature_dict.
      cat_feature_dict: Keys are context names and values are list of categorical
          features i.e. {"context_name" : [cat_0, ..., cat_k]}, where k is the
          number of categorical variables. If None, we use context names in the
          decomposition as the only categorical feature, i.e., k = 1.
      embs_feature_dict: Pre-trained continuous embedding features of each
          context.
      embs_dim_list: Embedding dimension for each categorical variable. The length
          equals the number of categorical features k. If None, the embedding
          dimension is set to 1 for each categorical variable.
      context_weight_dict: Known population weights of each context.
docstr_tok: ["Args",":","train_X",":","(","n","x","d",")","X","training","data",".","train_Y",":","(","n","x","1",")","Y","training","data",".","train_Yvar",":","(","n","x","1",")","Noise","variance","of","Y",".","If","None",",","we","use","an","inferred","noise","likelihood",".","decomposition",":","Keys","are","context","names",".","Values","are","the","indexes","of","parameters","belong","to","the","context",".","train_embedding",":","Whether","to","train","the","embedding","layer","or","not",".","If","False",",","the","model","will","use","pre-trained","embeddings","in","embs_feature_dict",".","cat_feature_dict",":","Keys","are","context","names","and","values","are","list","of","categorical","features","i.e",".","{","``","context_name","''",":","[","cat_0",",","...",",","cat_k","]","}",",","where","k","is","the","number","of","categorical","variables",".","If","None",",","we","use","context","names","in","the","decomposition","as","the","only","categorical","feature",",","i.e.",",","k","=","1.","embs_feature_dict",":","Pre-trained","continuous","embedding","features","of","each","context",".","embs_dim_list",":","Embedding","dimension","for","each","categorical","variable",".","The","length","equals","the","number","of","categorical","features","k.","If","None",",","the","embedding","dimension","is","set","to","1","for","each","categorical","variable",".","context_weight_dict",":","Known","population","weights","of","each","context","."]
code: def __init__(
      self,
      train_X: Tensor,
      train_Y: Tensor,
      train_Yvar: Optional[Tensor],
      decomposition: Dict[str, List[int]],
      train_embedding: bool = True,
      cat_feature_dict: Optional[Dict] = None,
      embs_feature_dict: Optional[Dict] = None,
      embs_dim_list: Optional[List[int]] = None,
      context_weight_dict: Optional[Dict] = None,
  ) -> None:
      r"""
      Args:
          train_X: (n x d) X training data.
          train_Y: (n x 1) Y training data.
          train_Yvar: (n x 1) Noise variance of Y. If None,
              we use an inferred noise likelihood.
          decomposition: Keys are context names. Values are the indexes of
              parameters belong to the context.
          train_embedding: Whether to train the embedding layer or not. If False,
              the model will use pre-trained embeddings in embs_feature_dict.
          cat_feature_dict: Keys are context names and values are list of categorical
              features i.e. {"context_name" : [cat_0, ..., cat_k]}, where k is the
              number of categorical variables. If None, we use context names in the
              decomposition as the only categorical feature, i.e., k = 1.
          embs_feature_dict: Pre-trained continuous embedding features of each
              context.
          embs_dim_list: Embedding dimension for each categorical variable. The length
              equals the number of categorical features k. If None, the embedding
              dimension is set to 1 for each categorical variable.
          context_weight_dict: Known population weights of each context.
      """
      super().__init__(
          train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar
      )
      self.covar_module = LCEAKernel(
          decomposition=decomposition,
          batch_shape=self._aug_batch_shape,
          train_embedding=train_embedding,
          cat_feature_dict=cat_feature_dict,
          embs_feature_dict=embs_feature_dict,
          embs_dim_list=embs_dim_list,
          context_weight_dict=context_weight_dict,
          device=train_X.device,
      )
      self.decomposition = decomposition
      self.to(train_X)
  
file_path: botorch/botorch/models/contextual.py
imports: from typing import Any, Dict, List, Optional
  from botorch.models.gp_regression import SingleTaskGP
  from botorch.models.kernels.contextual_lcea import LCEAKernel
  from botorch.models.kernels.contextual_sac import SACKernel
  from botorch.utils.datasets import SupervisedDataset
  from torch import Tensor
label_desc: Use image node_id 1 to create a new LCEAGP object from inherited base classes: SingleTaskGP with example: obj = LCEAGP(train_X, train_Y, train_Yvar, decomposition, train_embedding, cat_feature_dict, embs_feature_dict, embs_dim_list, context_weight_dict)