class_docstr: Sparse Multi-head Attention Module
class_docstr_tok: ["Sparse","Multi-head","Attention","Module"]
code: def forward(self, A, h):
      N = len(h)
      q = self.q_proj(h).reshape(N, self.head_dim, self.num_heads)
      q *= self.scaling
      k = self.k_proj(h).reshape(N, self.head_dim, self.num_heads)
      v = self.v_proj(h).reshape(N, self.head_dim, self.num_heads)
  
      ######################################################################
      # (HIGHLIGHT) Compute the multi-head attention with Sparse Matrix API
      ######################################################################
      attn = dglsp.bsddmm(A, q, k.transpose(1, 0))  # [N, N, nh]
      attn = attn.softmax()
      out = dglsp.bspmm(attn, v)
  
      return self.out_proj(out.reshape(N, -1))
  
imports: import dgl
  import dgl.nn
  import dgl.sparse
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from dgl.data import AsGraphPredDataset
  from dgl.dataloading import GraphDataLoader
  from ogb.graphproppred import collate_dgl, DglGraphPropPredDataset, Evaluator
  from ogb.graphproppred.mol_encoder import AtomEncoder
  from tqdm import tqdm
label_desc: Use image node_id 2 for calling the SparseMHA obj's underlying member method code with example usage: obj.forward(A, h) and returns: self