class_docstr: Parameters
  ----------
  supps_size: The encoding dimension of the support series.
  q_size: The encoding dimension of the query series.
  embed_size: The embedding dimension for keys, queries and values.
  num_heads: The number of attention heads
class_docstr_tok: ["Parameters","--","--","--","--","--","supps_size",":","The","encoding","dimension","of","the","support","series",".","q_size",":","The","encoding","dimension","of","the","query","series",".","embed_size",":","The","embedding","dimension","for","keys",",","queries","and","values",".","num_heads",":","The","number","of","attention","heads"]
docstr: Returns
  -------
  torch.Tensor:
      Aggregated values of size [batch, embed_size]
  list:
      optional, list of attention score tensors. Each tensor has size [n_supps, 1, support ts length].
      support ts length might be different from item to item in the list.
docstr_tok: ["Returns","--","--","--","-","torch.Tensor",":","Aggregated","values","of","size","[","batch",",","embed_size","]","list",":","optional",",","list","of","attention","score","tensors",".","Each","tensor","has","size","[","n_supps",",","1",",","support","ts","length","]",".","support","ts","length","might","be","different","from","item","to","item","in","the","list","."]
code: def forward(self, supps: SeriesBatch, query: SeriesBatch, mask=None):
      """
      Returns
      -------
      torch.Tensor:
          Aggregated values of size [batch, embed_size]
      list:
          optional, list of attention score tensors. Each tensor has size [n_supps, 1, support ts length].
          support ts length might be different from item to item in the list.
      """
      n_supps = supps.split_sections[0]
      pad_size = supps.sequences.size()[1]
      support_set = supps.sequences.reshape(
          -1, n_supps * pad_size, self.supps_embed_size
      )
      device = supps.sequences.device
      mask = torch.arange(pad_size, device=device)[
          None, :
      ] < supps.lengths[:, None].to(device)
      mask = ~mask.reshape(-1, n_supps * pad_size)
  
      values, attention = self.multihead_attn(
          query=query.sequences.unsqueeze(1),
          key=support_set,
          value=support_set,
          key_padding_mask=mask,
      )
      return values.squeeze(1), attention
  
file_path: gluonts/src/gluonts/nursery/few_shot_prediction/src/meta/models/EcDc/components/attention.py
imports: from typing import Union, Tuple
  from abc import ABC, abstractmethod
  import torch
  from torch import nn
  from meta.data.batch import SeriesBatch
label_desc: Use image node_id 2 for calling the MultiHeadSupportSetQueryAttention obj's underlying member method code with example usage: obj.forward(supps, query, mask) and returns: values,  attention