code: def graph_att_head(M, N, weight, ax, title):
      "credit: Jinjing Zhou"
      in_nodes = len(M)
      out_nodes = len(N)
  
      g = nx.bipartite.generators.complete_bipartite_graph(
          in_nodes, out_nodes
      )
      X, Y = bipartite.sets(g)
      height_in = 10
      height_out = height_in
      height_in_y = np.linspace(0, height_in, in_nodes)
      height_out_y = np.linspace(
          (height_in - height_out) / 2, height_out, out_nodes
      )
      pos = dict()
      pos.update(
          (n, (1, i)) for i, n in zip(height_in_y, X)
      )  # put nodes from X at x=1
      pos.update(
          (n, (3, i)) for i, n in zip(height_out_y, Y)
      )  # put nodes from Y at x=2
      ax.axis("off")
      ax.set_xlim(-1, 4)
      ax.set_title(title)
      nx.draw_networkx_nodes(
          g,
          pos,
          nodelist=range(in_nodes),
          node_color="r",
          node_size=50,
          ax=ax,
      )
      nx.draw_networkx_nodes(
          g,
          pos,
          nodelist=range(in_nodes, in_nodes + out_nodes),
          node_color="b",
          node_size=50,
          ax=ax,
      )
      for edge in g.edges():
          nx.draw_networkx_edges(
              g,
              pos,
              edgelist=[edge],
              width=weight[edge[0], edge[1] - in_nodes] * 1.5,
              ax=ax,
          )
      nx.draw_networkx_labels(
          g,
          pos,
          {i: label + "  " for i, label in enumerate(M)},
          horizontalalignment="right",
          font_size=8,
          ax=ax,
      )
      nx.draw_networkx_labels(
          g,
          pos,
          {i + in_nodes: "  " + label for i, label in enumerate(N)},
          horizontalalignment="left",
          font_size=8,
          ax=ax,
      )
  
file_path: dgl/examples/pytorch/transformer/modules/viz.py
imports: import os
  import matplotlib
  import matplotlib.animation
  import matplotlib.pyplot
  import networkx
  import numpy
  import torch
  from networkx.algorithms import bipartite
  import networkx
  from matplotlib.patches import ConnectionStyle, FancyArrowPatch
  from networkx.utils import is_string_like
label_desc: Use image node_id 5 for calling a global function with example usage: graph_att_head(M, N, weight, ax, title) without return types