code: def __init__(
      self,
      input_representation_size: int,
      hidden_size: int,
      output_representation_size: int,
      num_heads: int,
      query_representation_summarizer: AbstractVarSizedElementReduce,
      use_value_layer: bool = False,
  ):
      super().__init__()
      self.__query_layer = query_representation_summarizer
      self.__key_layer = nn.Linear(
          input_representation_size, hidden_size, bias=False
      )
      assert (
          hidden_size % num_heads == 0
      ), "Hidden size must be divisible by the number of heads."
      self.__use_value_layer = use_value_layer
      if use_value_layer:
          self.__value_layer = nn.Linear(
              input_representation_size, hidden_size, bias=False
          )
          self.__output_layer = nn.Linear(
              hidden_size, output_representation_size, bias=False
          )
      else:
          self.__output_layer = nn.Linear(
              input_representation_size * num_heads,
              output_representation_size,
              bias=False,
          )
      self.__num_heads = num_heads
  
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 1 to create a new MultiheadSelfAttentionVarSizedElementReduce object from inherited base classes: AbstractVarSizedElementReduce with example: obj = MultiheadSelfAttentionVarSizedElementReduce(input_representation_size, hidden_size, output_representation_size, num_heads, query_representation_summarizer, use_value_layer)