class_docstr: Message passing combined with transformer architecture
  
  The layer combines message passing performed using GroverMPNEncoder and uses it
  to generate query, key and value for multi-headed Attention block.
  
  Parameters
  ----------
  atom_messages: bool
      True if encoding atom-messages else False.
  input_dim: int
      Dimension of input features
  num_heads: int
      Number of attention heads
  depth: int
      Number of hops in a message passing iteration
  undirected: bool
      If set to `True`, the graph is considered as an undirected graph.
class_docstr_tok: ["Message","passing","combined","with","transformer","architecture","The","layer","combines","message","passing","performed","using","GroverMPNEncoder","and","uses","it","to","generate","query",",","key","and","value","for","multi-headed","Attention","block",".","Parameters","--","--","--","--","--","atom_messages",":","bool","True","if","encoding","atom-messages","else","False",".","input_dim",":","int","Dimension","of","input","features","num_heads",":","int","Number","of","attention","heads","depth",":","int","Number","of","hops","in","a","message","passing","iteration","undirected",":","bool","If","set","to","`","True","`",",","the","graph","is","considered","as","an","undirected","graph","."]
code: def __init__(
      self,
      atom_messages: bool,
      input_dim: int,
      num_heads: int,
      depth: int,
      undirected: bool = False,
      hidden_size: int = 128,
      dropout: float = 0.0,
      bias: bool = True,
      res_connection: bool = True,
      activation: str = "relu",
  ):
      super(GroverMTBlock, self).__init__()
      self.hidden_size = hidden_size
      self.atom_messages = atom_messages
      self.res_connection = res_connection
      self.num_heads = num_heads
      if activation == "relu":
          self.act_func = nn.ReLU()
      else:
          raise ValueError("Only relu activation is supported")
      self.dropout_layer = nn.Dropout(p=dropout)
  
      # Note: Elementwise affine has to be consistent with the pre-training phase
      self.layernorm = nn.LayerNorm(
          hidden_size, elementwise_affine=True
      )
      self.attn = torch.nn.MultiheadAttention(
          embed_dim=hidden_size,
          num_heads=num_heads,
          dropout=dropout,
          bias=bias,
          batch_first=True,
      )
  
      self.W_i = nn.Linear(input_dim, hidden_size, bias)
      self.W_o = nn.Linear(hidden_size * num_heads, hidden_size, bias)
      self.sublayer = SublayerConnection(hidden_size, dropout)
      self.attention_heads = nn.ModuleList()
      for i in range(num_heads):
          self.attention_heads.append(
              GroverAttentionHead(
                  hidden_size=hidden_size,
                  bias=bias,
                  dropout=dropout,
                  depth=depth,
                  atom_messages=atom_messages,
                  undirected=undirected,
              )
          )
  
imports: from typing import List, Dict
  import numpy
  from scipy import stats
  from deepchem.models.torch_models.readout import GroverReadout
  from deepchem.models.torch_models.layers import SublayerConnection, PositionwiseFeedForward
label_desc: Use image node_id 1 to create a new GroverMTBlock object from inherited base classes: nn with example: obj = GroverMTBlock(atom_messages, input_dim, num_heads, depth, undirected, hidden_size, dropout, bias, res_connection, activation)