code: def forward(
      self, inputs: ElementsToSummaryRepresentationInput
  ) -> torch.Tensor:
      queries = self.__query_layer(inputs)  # [num_samples, H]
      queries_per_element = queries[
          inputs.element_to_sample_map
      ]  # [num_elements, H]
      queries_per_element = queries_per_element.reshape(
          (
              queries_per_element.shape[0],
              self.__num_heads,
              queries_per_element.shape[1] // self.__num_heads,
          )
      )
  
      keys = self.__key_layer(
          inputs.element_embeddings
      )  # [num_elements, H]
      keys = keys.reshape(
          (
              keys.shape[0],
              self.__num_heads,
              keys.shape[1] // self.__num_heads,
          )
      )
  
      attention_scores = torch.einsum(
          "bhk,bhk->bh", queries_per_element, keys
      ) / sqrt(
          keys.shape[-1]
      )  # [num_elements, num_heads]
      attention_probs = torch.exp(
          scatter_log_softmax(
              attention_scores,
              index=inputs.element_to_sample_map,
              dim=0,
              eps=0,
          )
      )  # [num_elements, num_heads]
  
      if self.__use_value_layer:
          values = self.__value_layer(
              inputs.element_embeddings
          )  # [num_elements, hidden_size]
          values = values.reshape(
              (
                  values.shape[0],
                  self.__num_heads,
                  values.shape[1] // self.__num_heads,
              )
          )
          outputs = attention_probs.unsqueeze(-1) * values
      else:
          outputs = attention_probs.unsqueeze(
              -1
          ) * inputs.element_embeddings.unsqueeze(
              1
          )  # [num_elements, num_heads, D']
  
      outputs = outputs.reshape(
          (outputs.shape[0], -1)
      )  # [num_elements, num_heads * D']
  
      per_sample_outputs = scatter_sum(
          outputs,
          index=inputs.element_to_sample_map,
          dim=0,
          dim_size=inputs.num_samples,
      )  # [num_samples, num_heads, D']
  
      return self.__output_layer(per_sample_outputs)
  
imports: from typing_extensions import Literal
  import torch
  from abc import abstractmethod
  from math import sqrt
  from torch import nn
  from torch_scatter import scatter, scatter_log_softmax, scatter_sum
  from typing import NamedTuple, Union
label_desc: Use image node_id 2 for calling the MultiheadSelfAttentionVarSizedElementReduce obj's underlying member method code with example usage: obj.forward(inputs) and returns: self