class_docstr: Modular GNN which allows for easy swapping of GNN layers.
  
  Parameters
  ----------
  gnn_type: str
      The type of GNN layer to use. Must be one of "gin", "gcn", "graphsage", or "gat".
  num_layer: int
      The number of GNN layers to use.
  emb_dim: int
      The dimensionality of the node embeddings.
  num_tasks: int
      The number of tasks.
  graph_pooling: str
      The type of graph pooling to use. Must be one of "sum", "mean", "max", "attention" or "set2set".
      "sum" may cause issues with positive prediction loss.
  dropout: float, optional (default 0)
      The dropout probability.
  jump_knowledge: str, optional (default "last")
      The type of jump knowledge to use. [1] Must be one of "last", "sum", "max", or "concat".
      "last": Use the node representation from the last GNN layer.
      "concat": Concatenate the node representations from all GNN layers. This will increase the dimensionality of the node representations by a factor of `num_layer`.
      "max": Take the element-wise maximum of the node representations from all GNN layers.
      "sum": Take the element-wise sum of the node representations from all GNN layers. This may cause issues with positive prediction loss.
  task: str, optional (default "regression")
      The type of task.
      Unsupervised tasks:
      edge_pred: Edge prediction. Predicts whether an edge exists between two nodes.
      mask_nodes: Masking nodes. Predicts the masked node.
      mask_edges: Masking edges. Predicts the masked edge.
      infomax: Infomax. Maximizes mutual information between local node representations and a pooled global graph representation.
      context_pred: Context prediction. Predicts the surrounding context of a node.
      Supervised tasks:
      "regression" or "classification".
  mask_rate: float, optional (default 0.1)
      The rate at which to mask nodes or edges for mask_nodes and mask_edges tasks.
  mask_edge: bool, optional (default True)
      Whether to also mask connecting edges for mask_nodes tasks.
  context_size: int, optional (default 1)
      The size of the context to use for context prediction tasks.
  neighborhood_size: int, optional (default 3)
      The size of the neighborhood to use for context prediction tasks.
  context_mode: str, optional (default "cbow")
      The context mode to use for context prediction tasks. Must be one of "cbow" or "skipgram".
  neg_samples: int, optional (default 1)
      The number of negative samples to use for context prediction.
  
  Examples
  --------
  >>> import numpy as np
  >>> import deepchem as dc
  >>> from deepchem.feat.molecule_featurizers import SNAPFeaturizer
  >>> from deepchem.models.torch_models.gnn import GNNModular
  >>> featurizer = SNAPFeaturizer()
  >>> smiles = ["C1=CC=CC=C1", "C1=CC=CC=C1C=O", "C1=CC=CC=C1C(=O)O"]
  >>> features = featurizer.featurize(smiles)
  >>> dataset = dc.data.NumpyDataset(features, np.zeros(len(features)))
  >>> model = GNNModular(task="edge_pred")
  >>> loss = model.fit(dataset, nb_epoch=1)
  
  References
  ----------
  .. [1] Xu, K. et al. Representation Learning on Graphs with Jumping Knowledge Networks. Preprint at https://doi.org/10.48550/arXiv.1806.03536 (2018).
  .. [2] Hu, W. et al. Strategies for Pre-training Graph Neural Networks. Preprint at https://doi.org/10.48550/arXiv.1905.12265 (2020).
class_docstr_tok: ["Modular","GNN","which","allows","for","easy","swapping","of","GNN","layers",".","Parameters","--","--","--","--","--","gnn_type",":","str","The","type","of","GNN","layer","to","use",".","Must","be","one","of","``","gin","''",",","``","gcn","''",",","``","graphsage","''",",","or","``","gat","''",".","num_layer",":","int","The","number","of","GNN","layers","to","use",".","emb_dim",":","int","The","dimensionality","of","the","node","embeddings",".","num_tasks",":","int","The","number","of","tasks",".","graph_pooling",":","str","The","type","of","graph","pooling","to","use",".","Must","be","one","of","``","sum","''",",","``","mean","''",",","``","max","''",",","``","attention","''","or","``","set2set","''",".","``","sum","''","may","cause","issues","with","positive","prediction","loss",".","dropout",":","float",",","optional","(","default","0",")","The","dropout","probability",".","jump_knowledge",":","str",",","optional","(","default","``","last","''",")","The","type","of","jump","knowledge","to","use",".","[","1","]","Must","be","one","of","``","last","''",",","``","sum","''",",","``","max","''",",","or","``","concat","''",".","``","last","''",":","Use","the","node","representation","from","the","last","GNN","layer",".","``","concat","''",":","Concatenate","the","node","representations","from","all","GNN","layers",".","This","will","increase","the","dimensionality","of","the","node","representations","by","a","factor","of","`","num_layer","`",".","``","max","''",":","Take","the","element-wise","maximum","of","the","node","representations","from","all","GNN","layers",".","``","sum","''",":","Take","the","element-wise","sum","of","the","node","representations","from","all","GNN","layers",".","This","may","cause","issues","with","positive","prediction","loss",".","task",":","str",",","optional","(","default","``","regression","''",")","The","type","of","task",".","Unsupervised","tasks",":","edge_pred",":","Edge","prediction",".","Predicts","whether","an","edge","exists","between","two","nodes",".","mask_nodes",":","Masking","nodes",".","Predicts","the","masked","node",".","mask_edges",":","Masking","edges",".","Predicts","the","masked","edge",".","infomax",":","Infomax",".","Maximizes","mutual","information","between","local","node","representations","and","a","pooled","global","graph","representation",".","context_pred",":","Context","prediction",".","Predicts","the","surrounding","context","of","a","node",".","Supervised","tasks",":","``","regression","''","or","``","classification","''",".","mask_rate",":","float",",","optional","(","default","0.1",")","The","rate","at","which","to","mask","nodes","or","edges","for","mask_nodes","and","mask_edges","tasks",".","mask_edge",":","bool",",","optional","(","default","True",")","Whether","to","also","mask","connecting","edges","for","mask_nodes","tasks",".","context_size",":","int",",","optional","(","default","1",")","The","size","of","the","context","to","use","for","context","prediction","tasks",".","neighborhood_size",":","int",",","optional","(","default","3",")","The","size","of","the","neighborhood","to","use","for","context","prediction","tasks",".","context_mode",":","str",",","optional","(","default","``","cbow","''",")","The","context","mode","to","use","for","context","prediction","tasks",".","Must","be","one","of","``","cbow","''","or","``","skipgram","''",".","neg_samples",":","int",",","optional","(","default","1",")","The","number","of","negative","samples","to","use","for","context","prediction",".","Examples","--","--","--","--",">",">",">","import","numpy","as","np",">",">",">","import","deepchem","as","dc",">",">",">","from","deepchem.feat.molecule_featurizers","import","SNAPFeaturizer",">",">",">","from","deepchem.models.torch_models.gnn","import","GNNModular",">",">",">","featurizer","=","SNAPFeaturizer","(",")",">",">",">","smiles","=","[","``","C1=CC=CC=C1","''",",","``","C1=CC=CC=C1C=O","''",",","``","C1=CC=CC=C1C","(","=O",")","O","''","]",">",">",">","features","=","featurizer.featurize","(","smiles",")",">",">",">","dataset","=","dc.data.NumpyDataset","(","features",",","np.zeros","(","len","(","features",")",")",")",">",">",">","model","=","GNNModular","(","task=","''","edge_pred","''",")",">",">",">","loss","=","model.fit","(","dataset",",","nb_epoch=1",")","References","--","--","--","--","--","..","[","1","]","Xu",",","K.","et","al",".","Representation","Learning","on","Graphs","with","Jumping","Knowledge","Networks",".","Preprint","at","https",":","\/\/doi.org\/10.48550\/arXiv.1806.03536","(","2018",")",".","..","[","2","]","Hu",",","W.","et","al",".","Strategies","for","Pre-training","Graph","Neural","Networks",".","Preprint","at","https",":","\/\/doi.org\/10.48550\/arXiv.1905.12265","(","2020",")","."]
docstr: Builds the components of the GNNModular model. It initializes the encoders, batch normalization layers, pooling layers, and head layers based on the provided configuration. The method returns a dictionary containing the following components:
  
  Components list, type and description:
  --------------------------------------
  node_type_embedding: torch.nn.Embedding, an embedding layer for node types.
  
  chirality_embedding: torch.nn.Embedding, an embedding layer for chirality tags.
  
  gconvs: torch_geometric.nn.conv.MessagePassing, a list of graph convolutional layers (encoders) based on the specified GNN type (GIN, GCN, or GAT).
  
  batch_norms: torch.nn.BatchNorm1d, a list of batch normalization layers corresponding to the encoders.
  
  pool: Union[function,torch_geometric.nn.aggr.Aggregation], a pooling layer based on the specified graph pooling type (sum, mean, max, attention, or set2set).
  
  head: nn.Linear, a linear layer for the head of the model.
  
  These components are then used to construct the GNN and GNN_head modules for the GNNModular model.
docstr_tok: ["Builds","the","components","of","the","GNNModular","model",".","It","initializes","the","encoders",",","batch","normalization","layers",",","pooling","layers",",","and","head","layers","based","on","the","provided","configuration",".","The","method","returns","a","dictionary","containing","the","following","components",":","Components","list",",","type","and","description",":","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","--","node_type_embedding",":","torch.nn.Embedding",",","an","embedding","layer","for","node","types",".","chirality_embedding",":","torch.nn.Embedding",",","an","embedding","layer","for","chirality","tags",".","gconvs",":","torch_geometric.nn.conv.MessagePassing",",","a","list","of","graph","convolutional","layers","(","encoders",")","based","on","the","specified","GNN","type","(","GIN",",","GCN",",","or","GAT",")",".","batch_norms",":","torch.nn.BatchNorm1d",",","a","list","of","batch","normalization","layers","corresponding","to","the","encoders",".","pool",":","Union","[","function",",","torch_geometric.nn.aggr.Aggregation","]",",","a","pooling","layer","based","on","the","specified","graph","pooling","type","(","sum",",","mean",",","max",",","attention",",","or","set2set",")",".","head",":","nn.Linear",",","a","linear","layer","for","the","head","of","the","model",".","These","components","are","then","used","to","construct","the","GNN","and","GNN_head","modules","for","the","GNNModular","model","."]
code: def build_components(self):
      """
      Builds the components of the GNNModular model. It initializes the encoders, batch normalization layers, pooling layers, and head layers based on the provided configuration. The method returns a dictionary containing the following components:
  
      Components list, type and description:
      --------------------------------------
      node_type_embedding: torch.nn.Embedding, an embedding layer for node types.
  
      chirality_embedding: torch.nn.Embedding, an embedding layer for chirality tags.
  
      gconvs: torch_geometric.nn.conv.MessagePassing, a list of graph convolutional layers (encoders) based on the specified GNN type (GIN, GCN, or GAT).
  
      batch_norms: torch.nn.BatchNorm1d, a list of batch normalization layers corresponding to the encoders.
  
      pool: Union[function,torch_geometric.nn.aggr.Aggregation], a pooling layer based on the specified graph pooling type (sum, mean, max, attention, or set2set).
  
      head: nn.Linear, a linear layer for the head of the model.
  
      These components are then used to construct the GNN and GNN_head modules for the GNNModular model.
      """
  
      encoders, batch_norms = self.build_gnn(self.num_layer)
      components = {
          "node_type_embedding": torch.nn.Embedding(
              num_node_type, self.emb_dim
          ),
          "chirality_embedding": torch.nn.Embedding(
              num_chirality_tag, self.emb_dim
          ),
          "gconvs": encoders,
          "batch_norms": batch_norms,
      }
      self.gnn = GNN(
          components["node_type_embedding"],
          components["chirality_embedding"],
          components["gconvs"],
          components["batch_norms"],
          self.dropout,
          self.jump_knowledge,
      )
  
      if self.task in ("mask_nodes", "mask_edges"):
          linear_pred_nodes = torch.nn.Linear(
              self.emb_dim, num_node_type - 1
          )  # -1 to remove mask token
          linear_pred_edges = torch.nn.Linear(
              self.emb_dim, num_edge_type - 1
          )  # -1 to remove mask token
          components.update(
              {
                  "linear_pred_nodes": linear_pred_nodes,
                  "linear_pred_edges": linear_pred_edges,
              }
          )
  
      # for supervised tasks, add prediction head
      elif self.task in ("regression", "classification"):
          if self.graph_pooling == "sum":
              pool = global_add_pool
          elif self.graph_pooling == "mean":
              pool = global_mean_pool
          elif self.graph_pooling == "max":
              pool = global_max_pool
          elif self.graph_pooling == "attention":
              if self.jump_knowledge == "concat":
                  pool = AttentionalAggregation(
                      gate_nn=torch.nn.Linear(
                          (self.num_layer + 1) * self.emb_dim, 1
                      )
                  )
              else:
                  pool = AttentionalAggregation(
                      gate_nn=torch.nn.Linear(self.emb_dim, 1)
                  )
          elif self.graph_pooling == "set2set":
              set2setiter = 3
              if self.jump_knowledge == "concat":
                  pool = Set2Set(
                      (self.num_layer + 1) * self.emb_dim, set2setiter
                  )
              else:
                  pool = Set2Set(
                      self.emb_dim, processing_steps=set2setiter
                  )
  
          if self.graph_pooling == "set2set":
              mult = 2
          else:
              mult = 1
  
          if self.jump_knowledge == "concat":
              head = torch.nn.Linear(
                  mult * (self.num_layer + 1) * self.emb_dim,
                  self.output_dim,
              )
          else:
              head = torch.nn.Linear(
                  mult * self.emb_dim, self.output_dim
              )
  
          components.update({"pool": pool, "head": head})
  
          self.gnn_head = GNNHead(
              components["pool"],
              components["head"],
              self.task,
              self.num_tasks,
              self.num_classes,
          )
  
      elif self.task == "infomax":
          descrim = LocalGlobalDiscriminator(self.emb_dim)
          components.update(
              {"discriminator": descrim, "pool": global_mean_pool}
          )
  
      elif self.task == "context_pred":
          if self.graph_pooling == "sum":
              pool = global_add_pool
          elif self.graph_pooling == "mean":
              pool = global_mean_pool
          elif self.graph_pooling == "max":
              pool = global_max_pool
          elif self.graph_pooling == "attention":
              raise NotImplementedError(
                  "Attentional pooling not implemented for context prediction task."
              )
          elif self.graph_pooling == "set2set":
              raise NotImplementedError(
                  "Set2set pooling not implemented for context prediction task."
              )
  
          if (
              self.jump_knowledge == "concat"
          ):  # concat changes the emb_dim
              c_gconvs, c_batch_norms = self.build_gnn(self.num_layer)
          else:
              c_gconvs, c_batch_norms = self.build_gnn(
                  self.neighborhood_size - self.context_size
              )
          context_gnn_components = {
              "c_node_type_embedding": torch.nn.Embedding(
                  num_node_type, self.emb_dim
              ),
              "c_chirality_embedding": torch.nn.Embedding(
                  num_chirality_tag, self.emb_dim
              ),
              "c_gconvs": c_gconvs,
              "c_batch_norms": c_batch_norms,
          }
  
          self.context_gnn = GNN(
              context_gnn_components["c_node_type_embedding"],
              context_gnn_components["c_chirality_embedding"],
              context_gnn_components["c_gconvs"],
              context_gnn_components["c_batch_norms"],
              self.dropout,
              self.jump_knowledge,
          )
          components.update({"pool": pool, **context_gnn_components})
  
      return components
  
imports: import random
  import copy
  import torch
  import numpy
  from torch_geometric.nn import GINEConv, GCNConv, GATConv, SAGEConv, global_add_pool, global_mean_pool, global_max_pool
  from torch_geometric.nn.aggr import AttentionalAggregation, Set2Set
  from torch_geometric.nn.inits import uniform
  import torch.nn
  from torch.functional import F
  from deepchem.data import Dataset
  from deepchem.models.losses import SoftmaxCrossEntropy, EdgePredictionLoss, GraphNodeMaskingLoss, GraphEdgeMaskingLoss, DeepGraphInfomaxLoss, GraphContextPredLoss
  from deepchem.models.torch_models import ModularTorchModel
  from deepchem.feat.graph_data import BatchGraphData, GraphData, shortest_path_length
  from typing import Iterable, List, Tuple
  from deepchem.metrics import to_one_hot
label_desc: Use image node_id 2 for calling the GNNModular obj's underlying member method code with example usage: obj.build_components() and returns: components