class_docstr: Message Passing Neural Network,
  
  Message Passing Neural Networks [1]_ treat graph convolutional
  operations as an instantiation of a more general message
  passing schem. Recall that message passing in a graph is when
  nodes in a graph send each other "messages" and update their
  internal state as a consequence of these messages.
  
  Ordering structures in this model are built according to [2]_
  
  References
  ----------
  .. [1] Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, George E. Dahl.
      "Neural Message Passing for Quantum Chemistry." ICML 2017.
  .. [2] Vinyals, Oriol, Samy Bengio, and Manjunath Kudlur. "Order matters:
      Sequence to sequence for sets." arXiv preprint arXiv:1511.06391 (2015).
class_docstr_tok: ["Message","Passing","Neural","Network",",","Message","Passing","Neural","Networks","[","1","]","_","treat","graph","convolutional","operations","as","an","instantiation","of","a","more","general","message","passing","schem",".","Recall","that","message","passing","in","a","graph","is","when","nodes","in","a","graph","send","each","other","``","messages","''","and","update","their","internal","state","as","a","consequence","of","these","messages",".","Ordering","structures","in","this","model","are","built","according","to","[","2","]","_","References","--","--","--","--","--","..","[","1","]","Justin","Gilmer",",","Samuel","S.","Schoenholz",",","Patrick","F.","Riley",",","Oriol","Vinyals",",","George","E.","Dahl",".","``","Neural","Message","Passing","for","Quantum","Chemistry",".","''","ICML","2017",".","..","[","2","]","Vinyals",",","Oriol",",","Samy","Bengio",",","and","Manjunath","Kudlur",".","``","Order","matters",":","Sequence","to","sequence","for","sets",".","''","arXiv","preprint","arXiv:1511.06391","(","2015",")","."]
docstr: Parameters
  ----------
  n_tasks: int
      Number of tasks
  n_atom_feat: int, optional
      Number of features per atom.
  n_pair_feat: int, optional
      Number of features per pair of atoms.
  n_hidden: int, optional
      Number of units(convolution depths) in corresponding hidden layer
  n_graph_feat: int, optional
      Number of output features for each molecule(graph)
  dropout: float
      the dropout probablity to use.
  n_classes: int
      the number of classes to predict (only used in classification mode)
  uncertainty: bool
      if True, include extra outputs and loss terms to enable the uncertainty
      in outputs to be predicted
docstr_tok: ["Parameters","--","--","--","--","--","n_tasks",":","int","Number","of","tasks","n_atom_feat",":","int",",","optional","Number","of","features","per","atom",".","n_pair_feat",":","int",",","optional","Number","of","features","per","pair","of","atoms",".","n_hidden",":","int",",","optional","Number","of","units","(","convolution","depths",")","in","corresponding","hidden","layer","n_graph_feat",":","int",",","optional","Number","of","output","features","for","each","molecule","(","graph",")","dropout",":","float","the","dropout","probablity","to","use",".","n_classes",":","int","the","number","of","classes","to","predict","(","only","used","in","classification","mode",")","uncertainty",":","bool","if","True",",","include","extra","outputs","and","loss","terms","to","enable","the","uncertainty","in","outputs","to","be","predicted"]
code: def __init__(
      self,
      n_tasks,
      n_atom_feat=70,
      n_pair_feat=8,
      n_hidden=100,
      T=5,
      M=10,
      mode="regression",
      dropout=0.0,
      n_classes=2,
      uncertainty=False,
      batch_size=100,
      **kwargs,
  ):
      """
      Parameters
      ----------
      n_tasks: int
          Number of tasks
      n_atom_feat: int, optional
          Number of features per atom.
      n_pair_feat: int, optional
          Number of features per pair of atoms.
      n_hidden: int, optional
          Number of units(convolution depths) in corresponding hidden layer
      n_graph_feat: int, optional
          Number of output features for each molecule(graph)
      dropout: float
          the dropout probablity to use.
      n_classes: int
          the number of classes to predict (only used in classification mode)
      uncertainty: bool
          if True, include extra outputs and loss terms to enable the uncertainty
          in outputs to be predicted
      """
      if mode not in ["classification", "regression"]:
          raise ValueError(
              "mode must be either 'classification' or 'regression'"
          )
      self.n_tasks = n_tasks
      self.n_atom_feat = n_atom_feat
      self.n_pair_feat = n_pair_feat
      self.n_hidden = n_hidden
      self.T = T
      self.M = M
      self.mode = mode
      self.n_classes = n_classes
      self.uncertainty = uncertainty
      if uncertainty:
          if mode != "regression":
              raise ValueError(
                  "Uncertainty is only supported in regression mode"
              )
          if dropout == 0.0:
              raise ValueError(
                  "Dropout must be included to predict uncertainty"
              )
  
      # Build the model.
  
      atom_features = Input(shape=(self.n_atom_feat,))
      pair_features = Input(shape=(self.n_pair_feat,))
      atom_split = Input(shape=tuple(), dtype=tf.int32)
      atom_to_pair = Input(shape=(2,), dtype=tf.int32)
      n_samples = Input(shape=tuple(), dtype=tf.int32)
  
      message_passing = layers.MessagePassing(
          self.T,
          message_fn="enn",
          update_fn="gru",
          n_hidden=self.n_hidden,
      )([atom_features, pair_features, atom_to_pair])
  
      atom_embeddings = Dense(self.n_hidden)(message_passing)
  
      mol_embeddings = layers.SetGather(
          self.M, batch_size, n_hidden=self.n_hidden
      )([atom_embeddings, atom_split])
  
      dense1 = Dense(2 * self.n_hidden, activation=tf.nn.relu)(
          mol_embeddings
      )
  
      n_tasks = self.n_tasks
      if self.mode == "classification":
          n_classes = self.n_classes
          logits = Reshape((n_tasks, n_classes))(
              Dense(n_tasks * n_classes)(dense1)
          )
          logits = TrimGraphOutput()([logits, n_samples])
          output = Softmax()(logits)
          outputs = [output, logits]
          output_types = ["prediction", "loss"]
          loss = SoftmaxCrossEntropy()
      else:
          output = Dense(n_tasks)(dense1)
          output = TrimGraphOutput()([output, n_samples])
          if self.uncertainty:
              log_var = Dense(n_tasks)(dense1)
              log_var = TrimGraphOutput()([log_var, n_samples])
              var = Activation(tf.exp)(log_var)
              outputs = [output, var, output, log_var]
              output_types = ["prediction", "variance", "loss", "loss"]
  
              def loss(outputs, labels, weights):
                  (
                      output,
                      labels,
                  ) = dc.models.losses._make_tf_shapes_consistent(
                      outputs[0], labels[0]
                  )
                  output, labels = dc.models.losses._ensure_float(
                      output, labels
                  )
                  losses = (
                      tf.square(output - labels) / tf.exp(outputs[1])
                      + outputs[1]
                  )
                  w = weights[0]
                  if len(w.shape) < len(losses.shape):
                      if tf.is_tensor(w):
                          shape = tuple(w.shape.as_list())
                      else:
                          shape = w.shape
                      shape = tuple(
                          -1 if x is None else x for x in shape
                      )
                      w = tf.reshape(
                          w,
                          shape
                          + (1,) * (len(losses.shape) - len(w.shape)),
                      )
                  return tf.reduce_mean(losses * w) + sum(
                      self.model.losses
                  )
  
          else:
              outputs = [output]
              output_types = ["prediction"]
              loss = L2Loss()
      model = tf.keras.Model(
          inputs=[
              atom_features,
              pair_features,
              atom_split,
              atom_to_pair,
              n_samples,
          ],
          outputs=outputs,
      )
      super(MPNNModel, self).__init__(
          model,
          loss,
          output_types=output_types,
          batch_size=batch_size,
          **kwargs,
      )
  
imports: from collections.abc import Sequence
  import deepchem
  import numpy
  import tensorflow
  from typing import List, Union, Tuple, Iterable, Dict, Optional
  from deepchem.utils.typing import OneOrMany, LossFn, ActivationFn
  from deepchem.data import Dataset, pad_features
  from deepchem.feat.mol_graphs import ConvMol
  from deepchem.metrics import to_one_hot
  from deepchem.models import KerasModel, layers
  from deepchem.models.losses import L2Loss, SoftmaxCrossEntropy, Loss
  from tensorflow.keras.layers import Input, Dense, Reshape, Softmax, Dropout, Activation, BatchNormalization
  import warnings
label_desc: Use image node_id 1 to create a new MPNNModel object from inherited base classes: KerasModel with example: obj = MPNNModel(n_tasks, n_atom_feat, n_pair_feat, n_hidden, T, M, mode, dropout, n_classes, uncertainty, batch_size)