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: Loss that maximizes mutual information between local node representations and a pooled global graph representation. The positive and negative scores represent the similarity between local node representations and global graph representations of simlar and dissimilar graphs, respectively.
  
  Parameters
  ----------
  inputs: BatchedGraphData
      BatchedGraphData object containing the node features, edge indices, and graph indices for the batch of graphs.
docstr_tok: ["Loss","that","maximizes","mutual","information","between","local","node","representations","and","a","pooled","global","graph","representation",".","The","positive","and","negative","scores","represent","the","similarity","between","local","node","representations","and","global","graph","representations","of","simlar","and","dissimilar","graphs",",","respectively",".","Parameters","--","--","--","--","--","inputs",":","BatchedGraphData","BatchedGraphData","object","containing","the","node","features",",","edge","indices",",","and","graph","indices","for","the","batch","of","graphs","."]
code: def infomax_loss_loader(self, inputs):
      """
      Loss that maximizes mutual information between local node representations and a pooled global graph representation. The positive and negative scores represent the similarity between local node representations and global graph representations of simlar and dissimilar graphs, respectively.
  
      Parameters
      ----------
      inputs: BatchedGraphData
          BatchedGraphData object containing the node features, edge indices, and graph indices for the batch of graphs.
      """
      node_emb, inputs = self.model(inputs)
      summary_emb = torch.sigmoid(
          self.components["pool"](node_emb, inputs.graph_index)
      )
      positive_expanded_summary_emb = summary_emb[inputs.graph_index]
  
      shifted_summary_emb = summary_emb[
          cycle_index(len(summary_emb), 1)
      ]
      negative_expanded_summary_emb = shifted_summary_emb[
          inputs.graph_index
      ]
  
      positive_score = self.components["discriminator"](
          node_emb, positive_expanded_summary_emb
      )
      negative_score = self.components["discriminator"](
          node_emb, negative_expanded_summary_emb
      )
  
      return self.infomax_loss(positive_score, negative_score)
  
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 10 for calling the GNNModular obj's underlying member method code with example usage: obj.infomax_loss_loader(inputs) and returns: self