code: def test_mtgnn():
      """
      Testing MTGNN block
      """
      gcn_true = True
      build_adj = True
      dropout = 0.3
      subgraph_size = 20
      gcn_depth = 2
      num_nodes = 207
      node_dim = 40
      dilation_exponential = 1
      conv_channels = 32
      residual_channels = 32
      skip_channels = 64
      end_channels = 128
      in_dim = 2
      seq_in_len = 12
      seq_out_len = 10
      layers = 3
      batch_size = 16
      propalpha = 0.05
      tanhalpha = 3
      num_split = 1
      num_edges = 10
      kernel_size = 7
      kernel_set = [2, 3, 6, 7]
  
      device = torch.device(
          "cuda" if torch.cuda.is_available() else "cpu"
      )
      edge_index = barabasi_albert_graph(num_nodes, num_edges).to(
          device
      )
      A_tilde = (
          torch.sparse_coo_tensor(
              edge_index,
              torch.ones(edge_index.size(1)).to(device),
              (num_nodes, num_nodes),
          ).to_dense()
      ).to(device)
      model = MTGNN(
          gcn_true=gcn_true,
          build_adj=build_adj,
          gcn_depth=gcn_depth,
          num_nodes=num_nodes,
          kernel_size=kernel_size,
          kernel_set=kernel_set,
          dropout=dropout,
          subgraph_size=subgraph_size,
          node_dim=node_dim,
          dilation_exponential=dilation_exponential,
          conv_channels=conv_channels,
          residual_channels=residual_channels,
          skip_channels=skip_channels,
          end_channels=end_channels,
          seq_length=seq_in_len,
          in_dim=in_dim,
          out_dim=seq_out_len,
          layers=layers,
          propalpha=propalpha,
          tanhalpha=tanhalpha,
          layer_norm_affline=True,
      ).to(device)
      xd = 8
      FE = torch.rand(num_nodes, xd).to(device)
      model2 = MTGNN(
          gcn_true=gcn_true,
          build_adj=build_adj,
          gcn_depth=gcn_depth,
          num_nodes=num_nodes,
          kernel_size=kernel_size,
          kernel_set=kernel_set,
          dropout=dropout,
          subgraph_size=subgraph_size,
          node_dim=node_dim,
          dilation_exponential=dilation_exponential,
          conv_channels=conv_channels,
          residual_channels=residual_channels,
          skip_channels=skip_channels,
          end_channels=end_channels,
          seq_length=seq_in_len,
          in_dim=in_dim,
          out_dim=seq_out_len,
          layers=layers,
          propalpha=propalpha,
          tanhalpha=tanhalpha,
          layer_norm_affline=True,
          xd=xd,
      ).to(device)
  
      model3 = MTGNN(
          gcn_true=gcn_true,
          build_adj=False,
          gcn_depth=gcn_depth,
          num_nodes=num_nodes,
          kernel_size=kernel_size,
          kernel_set=kernel_set,
          dropout=dropout,
          subgraph_size=subgraph_size,
          node_dim=node_dim,
          dilation_exponential=dilation_exponential,
          conv_channels=conv_channels,
          residual_channels=residual_channels,
          skip_channels=skip_channels,
          end_channels=end_channels,
          seq_length=seq_in_len,
          in_dim=in_dim,
          out_dim=seq_out_len,
          layers=layers,
          propalpha=propalpha,
          tanhalpha=tanhalpha,
          layer_norm_affline=True,
          xd=xd,
      ).to(device)
      model4 = MTGNN(
          gcn_true=False,
          build_adj=build_adj,
          gcn_depth=gcn_depth,
          num_nodes=num_nodes,
          kernel_size=kernel_size,
          kernel_set=kernel_set,
          dropout=dropout,
          subgraph_size=subgraph_size,
          node_dim=node_dim,
          dilation_exponential=2,
          conv_channels=conv_channels,
          residual_channels=residual_channels,
          skip_channels=skip_channels,
          end_channels=end_channels,
          seq_length=seq_in_len,
          in_dim=in_dim,
          out_dim=seq_out_len,
          layers=layers,
          propalpha=propalpha,
          tanhalpha=tanhalpha,
          layer_norm_affline=False,
      ).to(device)
      trainx = (
          2 * torch.rand(batch_size, seq_in_len, num_nodes, in_dim) - 1
      ).to(device)
      trainx = trainx.transpose(1, 3)
      perm = torch.randperm(num_nodes).to(device)
      num_sub = int(num_nodes / num_split)
      for j in range(num_split):
          if j != num_split - 1:
              id = perm[j * num_sub : (j + 1) * num_sub]
          else:
              id = perm[j * num_sub :]
          tx = trainx[:, :, id, :]
          output = model(tx, A_tilde, idx=id)
          output = output.transpose(1, 3)
          assert output.shape == (batch_size, 1, num_nodes, seq_out_len)
          output2 = model2(tx, A_tilde, FE=FE)
          output2 = output2.transpose(1, 3)
          assert output2.shape == (
              batch_size,
              1,
              num_nodes,
              seq_out_len,
          )
          output3 = model3(tx, A_tilde, FE=FE)
          output3 = output3.transpose(1, 3)
          assert output3.shape == (
              batch_size,
              1,
              num_nodes,
              seq_out_len,
          )
          output4 = model4(tx, A_tilde)
          output4 = output4.transpose(1, 3)
          assert output4.shape == (
              batch_size,
              1,
              num_nodes,
              seq_out_len,
          )
  
      seq_in_len = 24
      seq_out_len = 5
      model = MTGNN(
          gcn_true=gcn_true,
          build_adj=build_adj,
          gcn_depth=gcn_depth,
          num_nodes=num_nodes,
          kernel_size=kernel_size,
          kernel_set=kernel_set,
          dropout=dropout,
          subgraph_size=subgraph_size,
          node_dim=node_dim,
          dilation_exponential=dilation_exponential,
          conv_channels=conv_channels,
          residual_channels=residual_channels,
          skip_channels=skip_channels,
          end_channels=end_channels,
          seq_length=seq_in_len,
          in_dim=in_dim,
          out_dim=seq_out_len,
          layers=layers,
          propalpha=propalpha,
          tanhalpha=tanhalpha,
          layer_norm_affline=False,
      ).to(device)
      dilation_exponential = 2
      build_adj = False
      model2 = MTGNN(
          gcn_true=gcn_true,
          build_adj=build_adj,
          gcn_depth=gcn_depth,
          num_nodes=num_nodes,
          kernel_size=kernel_size,
          kernel_set=kernel_set,
          dropout=dropout,
          subgraph_size=subgraph_size,
          node_dim=node_dim,
          dilation_exponential=dilation_exponential,
          conv_channels=conv_channels,
          residual_channels=residual_channels,
          skip_channels=skip_channels,
          end_channels=end_channels,
          seq_length=seq_in_len,
          in_dim=in_dim,
          out_dim=seq_out_len,
          layers=layers,
          propalpha=propalpha,
          tanhalpha=tanhalpha,
          layer_norm_affline=True,
          xd=xd,
      ).to(device)
      model3 = MTGNN(
          gcn_true=False,
          build_adj=build_adj,
          gcn_depth=gcn_depth,
          num_nodes=num_nodes,
          kernel_size=kernel_size,
          kernel_set=kernel_set,
          dropout=dropout,
          subgraph_size=subgraph_size,
          node_dim=node_dim,
          dilation_exponential=dilation_exponential,
          conv_channels=conv_channels,
          residual_channels=residual_channels,
          skip_channels=skip_channels,
          end_channels=end_channels,
          seq_length=seq_in_len,
          in_dim=in_dim,
          out_dim=seq_out_len,
          layers=layers,
          propalpha=propalpha,
          tanhalpha=tanhalpha,
          layer_norm_affline=False,
      ).to(device)
      trainx = (
          2 * torch.rand(batch_size, seq_in_len, num_nodes, in_dim) - 1
      ).to(device)
      trainx = trainx.transpose(1, 3)
      for j in range(num_split):
          if j != num_split - 1:
              id = perm[j * num_sub : (j + 1) * num_sub]
          else:
              id = perm[j * num_sub :]
          tx = trainx[:, :, id, :]
          output = model(tx, A_tilde, idx=id)
          output = output.transpose(1, 3)
          assert output.shape == (batch_size, 1, num_nodes, seq_out_len)
          output2 = model2(tx, A_tilde, idx=id, FE=FE)
          output2 = output2.transpose(1, 3)
          assert output2.shape == (
              batch_size,
              1,
              num_nodes,
              seq_out_len,
          )
          output3 = model3(tx, A_tilde)
          output3 = output3.transpose(1, 3)
          assert output3.shape == (
              batch_size,
              1,
              num_nodes,
              seq_out_len,
          )
  
imports: import torch
  import math
  import numpy
  import networkx
  from torch_geometric.data import Data
  from torch_geometric.utils import barabasi_albert_graph
  from torch_geometric.transforms import LaplacianLambdaMax
  from torch_geometric_temporal.nn.attention import TemporalConv, STConv, ASTGCN, MSTGCN, MTGNN, ChebConvAttention, AAGCN, GraphAAGCN, DNNTSP
  from torch_geometric_temporal.nn.attention import GMAN, SpatioTemporalAttention, SpatioTemporalEmbedding
label_desc: Use image node_id 11 for calling a global function with example usage: test_mtgnn() without return types