class_docstr: MatErials Graph Network
  
  A model for predicting crystal and molecular properties using GraphNetworks.
  
  Example
  -------
  >>> import numpy as np
  >>> from torch_geometric.data import Batch
  >>> from deepchem.feat import GraphData
  >>> n_nodes, n_node_features = 5, 10
  >>> n_edges, n_edge_attrs = 5, 2
  >>> n_global_features = 4
  >>> node_features = np.random.randn(n_nodes, n_node_features)
  >>> edge_attrs = np.random.randn(n_edges, n_edge_attrs)
  >>> edge_index = np.array([[0, 1, 2, 3, 4], [1, 2, 3, 4, 0]], dtype=np.int64)
  >>> global_features = np.random.randn(1, n_global_features)
  >>> graph = GraphData(node_features, edge_index, edge_attrs, global_features=global_features)
  >>> batch = Batch()
  >>> batch = batch.from_data_list([graph.to_pyg_graph()])
  >>> model = MEGNet(n_node_features=n_node_features, n_edge_features=n_edge_attrs, n_global_features=n_global_features)
  >>> pred = model(batch)
  
  Note
  ----
  This class requires torch-geometric to be installed.
docstr: Parameters
  ----------
  n_node_features: int
      Number of features in a node
  n_edge_features: int
      Number of features in a edge
  n_global_features: int
      Number of global features
  n_blocks: int
      Number of GraphNetworks block to use in update
  is_undirected: bool, optional (default True)
      True when the graph is undirected graph , otherwise False
  residual_connection: bool, optional (default True)
      If True, the layer uses a residual connection during training
  n_tasks: int, default 1
      The number of tasks
  mode: str, default 'regression'
      The model type - classification or regression
  n_classes: int, default 2
      The number of classes to predict (used only in classification mode).
docstr_tok: ["Parameters","--","--","--","--","--","n_node_features",":","int","Number","of","features","in","a","node","n_edge_features",":","int","Number","of","features","in","a","edge","n_global_features",":","int","Number","of","global","features","n_blocks",":","int","Number","of","GraphNetworks","block","to","use","in","update","is_undirected",":","bool",",","optional","(","default","True",")","True","when","the","graph","is","undirected","graph",",","otherwise","False","residual_connection",":","bool",",","optional","(","default","True",")","If","True",",","the","layer","uses","a","residual","connection","during","training","n_tasks",":","int",",","default","1","The","number","of","tasks","mode",":","str",",","default","'regression'","The","model","type","-","classification","or","regression","n_classes",":","int",",","default","2","The","number","of","classes","to","predict","(","used","only","in","classification","mode",")","."]
code: def __init__(
      self,
      n_node_features: int = 32,
      n_edge_features: int = 32,
      n_global_features: int = 32,
      n_blocks: int = 1,
      is_undirected: bool = True,
      residual_connection: bool = True,
      mode: str = "regression",
      n_classes: int = 2,
      n_tasks: int = 1,
  ):
      """
  
      Parameters
      ----------
      n_node_features: int
          Number of features in a node
      n_edge_features: int
          Number of features in a edge
      n_global_features: int
          Number of global features
      n_blocks: int
          Number of GraphNetworks block to use in update
      is_undirected: bool, optional (default True)
          True when the graph is undirected graph , otherwise False
      residual_connection: bool, optional (default True)
          If True, the layer uses a residual connection during training
      n_tasks: int, default 1
          The number of tasks
      mode: str, default 'regression'
          The model type - classification or regression
      n_classes: int, default 2
          The number of classes to predict (used only in classification mode).
      """
      super(MEGNet, self).__init__()
      try:
          from torch_geometric.nn import Set2Set
      except ModuleNotFoundError:
          raise ImportError(
              "MEGNet model requires torch_geometric to be installed"
          )
  
      if mode not in ["classification", "regression"]:
          raise ValueError(
              "mode must be either 'classification' or 'regression'"
          )
      self.n_node_features = n_node_features
      self.n_edge_features = n_edge_features
      self.n_global_features = n_global_features
      self.megnet_blocks = nn.ModuleList()
      self.n_blocks = n_blocks
      for i in range(n_blocks):
          self.megnet_blocks.append(
              GN(
                  n_node_features=n_node_features,
                  n_edge_features=n_edge_features,
                  n_global_features=n_global_features,
                  is_undirected=is_undirected,
                  residual_connection=residual_connection,
              )
          )
      self.n_tasks = n_tasks
      self.mode = mode
      self.n_classes = n_classes
  
      self.set2set_nodes = Set2Set(
          in_channels=n_node_features, processing_steps=3, num_layers=1
      )
      self.set2set_edges = Set2Set(
          in_channels=n_edge_features, processing_steps=3, num_layers=1
      )
  
      self.dense = nn.Sequential(
          nn.Linear(
              in_features=2 * n_node_features
              + 2 * n_edge_features
              + n_global_features,
              out_features=32,
          ),
          nn.Linear(in_features=32, out_features=16),
      )
  
      if self.mode == "regression":
          self.out = nn.Linear(in_features=16, out_features=n_tasks)
      elif self.mode == "classification":
          self.out = nn.Linear(
              in_features=16, out_features=n_tasks * n_classes
          )
  
file_path: deepchem/deepchem/models/torch_models/megnet.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from deepchem.models.losses import Loss, L2Loss, SparseSoftmaxCrossEntropy
  from deepchem.models.torch_models.layers import GraphNetwork
  from deepchem.models.torch_models import TorchModel
label_desc: Use image node_id 1 to create a new MEGNet object from inherited base classes: nn with example: obj = MEGNet(n_node_features, n_edge_features, n_global_features, n_blocks, is_undirected, residual_connection, mode, n_classes, n_tasks)