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 forward(self, batch):
      f_atoms, f_bonds, a2b, b2a, b2revb, a_scope, b_scope, a2a = batch
      if self.atom_messages:
          if f_atoms.shape[1] != self.hidden_size:
              f_atoms = self.W_i(f_atoms)
              f_atoms = self.dropout_layer(
                  self.layernorm(self.act_func(f_atoms))
              )
  
      else:
          if f_bonds.shape[1] != self.hidden_size:
              f_bonds = self.W_i(f_bonds)
              f_bonds = self.dropout_layer(
                  self.layernorm(self.act_func(f_bonds))
              )
  
      queries, keys, values = [], [], []
      for head in self.attention_heads:
          q, k, v = head(
              f_atoms, f_bonds, a2b=a2b, b2a=b2a, b2revb=b2revb, a2a=a2a
          )
          queries.append(q.unsqueeze(1))
          keys.append(k.unsqueeze(1))
          values.append(v.unsqueeze(1))
      queries = torch.cat(queries, dim=1)
      keys = torch.cat(keys, dim=1)
      values = torch.cat(values, dim=1)
  
      # multi-headed attention
      x_out, _ = self.attn(queries, keys, values)
      x_out = x_out.reshape(x_out.shape[0], -1)
      x_out = self.W_o(x_out)
  
      # support no residual connection in MTBlock.
      if self.res_connection:
          if self.atom_messages:
              f_atoms = self.sublayer(f_atoms, x_out)
          else:
              f_bonds = self.sublayer(f_bonds, x_out)
      batch = f_atoms, f_bonds, a2b, b2a, b2revb, a_scope, b_scope, a2a
      return batch
  
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 2 for calling the GroverMTBlock obj's underlying member method code with example usage: obj.forward(batch) and returns: batch