class_docstr: GNN module for the GNNModular model.
  
  This module is responsible for the graph neural network layers in the GNNModular model.
  
  Parameters
  ----------
  node_type_embedding: torch.nn.Embedding
      Embedding layer for node types.
  chirality_embedding: torch.nn.Embedding
      Embedding layer for chirality tags.
  gconvs: torch.nn.ModuleList
      ModuleList of graph convolutional layers.
  batch_norms: torch.nn.ModuleList
      ModuleList of batch normalization layers.
  dropout: int
      Dropout probability.
  jump_knowledge: str
      The type of jump knowledge to use. [1] Must be one of "last", "sum", "max", "concat" or "none".
      "last": Use the node representation from the last GNN layer.
      "concat": Concatenate the node representations from all GNN layers.
      "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.
  init_emb: bool
      Whether to initialize the embedding layers with Xavier uniform initialization.
  
  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).
  
  Example
  -------
  >>> from deepchem.models.torch_models.gnn import GNNModular
  >>> from deepchem.feat.graph_data import BatchGraphData
  >>> from deepchem.feat.molecule_featurizers import SNAPFeaturizer
  >>> featurizer = SNAPFeaturizer()
  >>> smiles = ["C1=CC=CC=C1", "C1=CC=CC=C1C=O", "C1=CC=CC=C1C(=O)O"]
  >>> features = featurizer.featurize(smiles)
  >>> modular = GNNModular(emb_dim = 8, task = "edge_pred")
  >>> batched_graph = BatchGraphData(features).numpy_to_torch(device=modular.device)
  >>> gnnmodel = modular.gnn
  >>> print(gnnmodel(batched_graph)[0].shape)
  torch.Size([23, 8])
class_docstr_tok: ["GNN","module","for","the","GNNModular","model",".","This","module","is","responsible","for","the","graph","neural","network","layers","in","the","GNNModular","model",".","Parameters","--","--","--","--","--","node_type_embedding",":","torch.nn.Embedding","Embedding","layer","for","node","types",".","chirality_embedding",":","torch.nn.Embedding","Embedding","layer","for","chirality","tags",".","gconvs",":","torch.nn.ModuleList","ModuleList","of","graph","convolutional","layers",".","batch_norms",":","torch.nn.ModuleList","ModuleList","of","batch","normalization","layers",".","dropout",":","int","Dropout","probability",".","jump_knowledge",":","str","The","type","of","jump","knowledge","to","use",".","[","1","]","Must","be","one","of","``","last","''",",","``","sum","''",",","``","max","''",",","``","concat","''","or","``","none","''",".","``","last","''",":","Use","the","node","representation","from","the","last","GNN","layer",".","``","concat","''",":","Concatenate","the","node","representations","from","all","GNN","layers",".","``","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",".","init_emb",":","bool","Whether","to","initialize","the","embedding","layers","with","Xavier","uniform","initialization",".","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",")",".","Example","--","--","--","-",">",">",">","from","deepchem.models.torch_models.gnn","import","GNNModular",">",">",">","from","deepchem.feat.graph_data","import","BatchGraphData",">",">",">","from","deepchem.feat.molecule_featurizers","import","SNAPFeaturizer",">",">",">","featurizer","=","SNAPFeaturizer","(",")",">",">",">","smiles","=","[","``","C1=CC=CC=C1","''",",","``","C1=CC=CC=C1C=O","''",",","``","C1=CC=CC=C1C","(","=O",")","O","''","]",">",">",">","features","=","featurizer.featurize","(","smiles",")",">",">",">","modular","=","GNNModular","(","emb_dim","=","8",",","task","=","``","edge_pred","''",")",">",">",">","batched_graph","=","BatchGraphData","(","features",")",".numpy_to_torch","(","device=modular.device",")",">",">",">","gnnmodel","=","modular.gnn",">",">",">","print","(","gnnmodel","(","batched_graph",")","[","0","]",".shape",")","torch.Size","(","[","23",",","8","]",")"]
code: def __init__(
      self,
      node_type_embedding,
      chirality_embedding,
      gconvs,
      batch_norms,
      dropout,
      jump_knowledge,
      init_emb=False,
  ):
      super(GNN, self).__init__()
  
      self.node_type_embedding = node_type_embedding
      self.chirality_embedding = chirality_embedding
      self.gconv = gconvs
      self.batch_norms = batch_norms
      self.dropout = dropout
      self.num_layer = len(gconvs)
      self.jump_knowledge = jump_knowledge
  
      # may mess with loading pretrained weights
      if init_emb:
          torch.nn.init.xavier_uniform_(
              self.node_type_embedding.weight.data
          )
          torch.nn.init.xavier_uniform_(
              self.chirality_embedding.weight.data
          )
  
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 1 to create a new GNN object from inherited base classes: torch.nn with example: obj = GNN(node_type_embedding, chirality_embedding, gconvs, batch_norms, dropout, jump_knowledge, init_emb)