class_docstr: The Latent Context Embedding Additive (LCE-A) Kernel.
  
  This kernel is similar to the SACKernel, and is used when context breakdowns are
  unbserverable. It assumes the same additive structure and a spatial kernel shared
  across contexts. Rather than assuming independence, LCEAKernel models the
  correlation in the latent functions for each context through learning context
  embeddings.
docstr: Args:
      decomposition: Keys index context names. Values are the indexes of
          parameters belong to the context.
      batch_shape: Batch shape as usual for gpytorch kernels. Model does not
          support batch training. When batch_shape is non-empty, it is used for
          loading hyper-parameter values generated from MCMC sampling.
      train_embedding: A boolean indictor of whether to learn context embeddings.
      cat_feature_dict: Keys are context names and values are list of categorical
          features i.e. {"context_name" : [cat_0, ..., cat_k]}. k equals the
          number of categorical variables. If None, uses 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 to num 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",":","decomposition",":","Keys","index","context","names",".","Values","are","the","indexes","of","parameters","belong","to","the","context",".","batch_shape",":","Batch","shape","as","usual","for","gpytorch","kernels",".","Model","does","not","support","batch","training",".","When","batch_shape","is","non-empty",",","it","is","used","for","loading","hyper-parameter","values","generated","from","MCMC","sampling",".","train_embedding",":","A","boolean","indictor","of","whether","to","learn","context","embeddings",".","cat_feature_dict",":","Keys","are","context","names","and","values","are","list","of","categorical","features","i.e",".","{","``","context_name","''",":","[","cat_0",",","...",",","cat_k","]","}",".","k","equals","the","number","of","categorical","variables",".","If","None",",","uses","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","to","num","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,
      decomposition: Dict[str, List[int]],
      batch_shape: torch.Size,
      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,
      device: Optional[torch.device] = None,
  ) -> None:
      r"""
      Args:
          decomposition: Keys index context names. Values are the indexes of
              parameters belong to the context.
          batch_shape: Batch shape as usual for gpytorch kernels. Model does not
              support batch training. When batch_shape is non-empty, it is used for
              loading hyper-parameter values generated from MCMC sampling.
          train_embedding: A boolean indictor of whether to learn context embeddings.
          cat_feature_dict: Keys are context names and values are list of categorical
              features i.e. {"context_name" : [cat_0, ..., cat_k]}. k equals the
              number of categorical variables. If None, uses 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 to num 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__(batch_shape=batch_shape)
      self.batch_shape = batch_shape
      self.train_embedding = train_embedding
      self._device = device
  
      self.num_param = len(next(iter(decomposition.values())))
      self.context_list = list(decomposition.keys())
      self.num_contexts = len(self.context_list)
  
      # get parameter space decomposition
      for active_parameters in decomposition.values():
          # check number of parameters are same in each decomp
          if len(active_parameters) != self.num_param:
              raise ValueError(
                  "The number of parameters needs to be same across all contexts."
              )
      # reorder the parameter list based on decomposition such that
      # parameters for each context are contiguous and in the same order for each
      # context
      self.permutation = get_permutation(decomposition=decomposition)
      # get context features and set emb dim
      self.context_cat_feature = None
      self.context_emb_feature = None
      self.n_embs = 0
      self.emb_weight_matrix_list = None
      self.emb_dims = None
      self._set_context_features(
          cat_feature_dict=cat_feature_dict,
          embs_feature_dict=embs_feature_dict,
          embs_dim_list=embs_dim_list,
      )
      # contruct embedding layer
      if train_embedding:
          self._set_emb_layers()
      # task covariance matrix
      self.task_covar_module = MaternKernel(
          nu=2.5,
          ard_num_dims=self.n_embs,
          batch_shape=batch_shape,
          lengthscale_prior=GammaPrior(3.0, 6.0),
      )
      # base kernel
      self.base_kernel = MaternKernel(
          nu=2.5,
          ard_num_dims=self.num_param,
          batch_shape=batch_shape,
          lengthscale_prior=GammaPrior(3.0, 6.0),
      )
      # outputscales for each context (note this is like sqrt of outputscale)
      self.context_weight = None
      if context_weight_dict is None:
          outputscale_list = torch.zeros(
              *batch_shape, self.num_contexts, device=self.device
          )
      else:
          outputscale_list = torch.zeros(
              *batch_shape, 1, device=self.device
          )
          self.context_weight = torch.tensor(
              [context_weight_dict[c] for c in self.context_list],
              device=self.device,
          )
      self.register_parameter(
          name="raw_outputscale_list",
          parameter=torch.nn.Parameter(outputscale_list),
      )
      self.register_prior(
          "outputscale_list_prior",
          GammaPrior(2.0, 15.0),
          lambda m: m.outputscale_list,
          lambda m, v: m._set_outputscale_list(v),
      )
      self.register_constraint("raw_outputscale_list", Positive())
  
imports: from typing import Any, Dict, List, Optional
  import torch
  from gpytorch.constraints import Positive
  from gpytorch.kernels.kernel import Kernel
  from gpytorch.kernels.matern_kernel import MaternKernel
  from gpytorch.priors.torch_priors import GammaPrior
  from linear_operator.operators import DiagLinearOperator
  from linear_operator.operators.dense_linear_operator import DenseLinearOperator
  from torch import Tensor
  from torch.nn import ModuleList
label_desc: Use image node_id 1 to create a new LCEAKernel object from inherited base classes: Kernel with example: obj = LCEAKernel(decomposition, batch_shape, train_embedding, cat_feature_dict, embs_feature_dict, embs_dim_list, context_weight_dict, device)