Fixes

        yield (inputs, [y_b], [w_b])

#class GraphConvModel(KerasModel):

#

#  def __init__(self,

#               n_tasks,

#               graph_conv_layers=[64, 64],

#               dense_layer_size=128,

#               dropout=0.0,

#               mode="classification",

#               number_atom_features=75,

#               n_classes=2,

#               uncertainty=False,

#               batch_size=100,

#               **kwargs):

#    """

#    Parameters

#    ----------

#    n_tasks: int

#      Number of tasks

#    graph_conv_layers: list of int

#      Width of channels for the Graph Convolution Layers

#    dense_layer_size: int

#      Width of channels for Atom Level Dense Layer before GraphPool

#    dropout: list or float

#      the dropout probablity to use for each layer.  The length of this list should equal

#      len(graph_conv_layers)+1 (one value for each convolution layer, and one for the

#      dense layer).  Alternatively this may be a single value instead of a list, in which

#      case the same value is used for every layer.

#    mode: str

#      Either "classification" or "regression"

#    number_atom_features: int

#        75 is the default number of atom features created, but

#        this can vary if various options are passed to the

#        function atom_features in graph_features

#    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.mode = mode

#    self.dense_layer_size = dense_layer_size

#    self.graph_conv_layers = graph_conv_layers

#    self.number_atom_features = number_atom_features

#    self.n_classes = n_classes

#    self.uncertainty = uncertainty

#    if not isinstance(dropout, collections.Sequence):

#      dropout = [dropout] * (len(graph_conv_layers) + 1)

#    if len(dropout) != len(graph_conv_layers) + 1:

#      raise ValueError('Wrong number of dropout probabilities provided')

#    self.dropout = dropout

#    if uncertainty:

#      if mode != "regression":

#        raise ValueError("Uncertainty is only supported in regression mode")

#      if any(d == 0.0 for d in dropout):

#        raise ValueError(

#            'Dropout must be included in every layer to predict uncertainty')

#

#    # Build the model.

#

#    atom_features = Input(shape=(self.number_atom_features,))

#    degree_slice = Input(shape=(2,), dtype=tf.int32)

#    membership = Input(shape=tuple(), dtype=tf.int32)

#    n_samples = Input(shape=tuple(), dtype=tf.int32)

#    dropout_switch = tf.keras.Input(shape=tuple())

#

#    self.deg_adjs = []

#    for i in range(0, 10 + 1):

#      deg_adj = Input(shape=(i + 1,), dtype=tf.int32)

#      self.deg_adjs.append(deg_adj)

#    in_layer = atom_features

#    for layer_size, dropout in zip(self.graph_conv_layers, self.dropout):

#      gc1_in = [in_layer, degree_slice, membership] + self.deg_adjs

#      gc1 = layers.GraphConv(layer_size, activation_fn=tf.nn.relu)(gc1_in)

#      batch_norm1 = BatchNormalization(fused=False)(gc1)

#      if dropout > 0.0:

#        batch_norm1 = layers.SwitchedDropout(rate=dropout)(

#            [batch_norm1, dropout_switch])

#      gp_in = [batch_norm1, degree_slice, membership] + self.deg_adjs

#      in_layer = layers.GraphPool()(gp_in)

#    dense = Dense(self.dense_layer_size, activation=tf.nn.relu)(in_layer)

#    batch_norm3 = BatchNormalization(fused=False)(dense)

#    if self.dropout[-1] > 0.0:

#      batch_norm3 = layers.SwitchedDropout(rate=self.dropout[-1])(

#          [batch_norm3, dropout_switch])

#    self.neural_fingerprint = batch_norm3

#    self.neural_fingerprint = layers.GraphGather(

#        batch_size=batch_size,

#        activation_fn=tf.nn.tanh)([batch_norm3, degree_slice, membership] +

#                                  self.deg_adjs)

#

#    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)(

#          self.neural_fingerprint))

#      logits = TrimGraphOutput()([logits, n_samples])

#      output = Softmax()(logits)

#      outputs = [output, logits]

#      output_types = ['prediction', 'loss']

#      loss = SoftmaxCrossEntropy()

#    else:

#      output = Dense(n_tasks)(self.neural_fingerprint)

#      output = TrimGraphOutput()([output, n_samples])

#      if self.uncertainty:

#        log_var = Dense(n_tasks)(self.neural_fingerprint)

#        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):

#          diff = labels[0] - outputs[0]

#          return tf.reduce_mean(diff * diff / tf.exp(outputs[1]) + outputs[1])

#      else:

#        outputs = [output]

#        output_types = ['prediction']

#        loss = L2Loss()

#    model = tf.keras.Model(

#        inputs=[

#            atom_features, degree_slice, membership, n_samples, dropout_switch

#        ] + self.deg_adjs,

#        outputs=outputs)

#    super(GraphConvModel, self).__init__(

#        model, loss, output_types=output_types, batch_size=batch_size, **kwargs)

#

#  def default_generator(self,

#                        dataset,

#                        epochs=1,

#                        mode='fit',

#                        deterministic=True,

#                        pad_batches=True):

#    for epoch in range(epochs):

#      for (X_b, y_b, w_b, ids_b) in dataset.iterbatches(

#          batch_size=self.batch_size,

#          deterministic=deterministic,

#          pad_batches=pad_batches):

#        if self.mode == 'classification':

#          y_b = to_one_hot(y_b.flatten(), self.n_classes).reshape(

#              -1, self.n_tasks, self.n_classes)

#        multiConvMol = ConvMol.agglomerate_mols(X_b)

#        n_samples = np.array(X_b.shape[0])

#        if mode == 'predict':

#          dropout = np.array(0.0)

#        else:

#          dropout = np.array(1.0)

#        inputs = [

#            multiConvMol.get_atom_features(), multiConvMol.deg_slice,

#            np.array(multiConvMol.membership), n_samples, dropout

#        ]

#        for i in range(1, len(multiConvMol.get_deg_adjacency_lists())):

#          inputs.append(multiConvMol.get_deg_adjacency_lists()[i])

#        yield (inputs, [y_b], [w_b])

class MPNNModel(KerasModel):

  """ Message Passing Neural Network,

      default structures built according to https://arxiv.org/abs/1511.06391 """

          for i, t in enumerate(var):

            variances[i].append(t)

      if embedding:

        ##################################

        print("len(self._embedding_outputs)")

        print(len(self._embedding_outputs))

        ##################################

        embeddings = [output_values[i] for i in self._embedding_outputs]

        ##################################

        print("embeddings[0].shape")

        print(embeddings[0].shape)

        ##################################

      if self._prediction_outputs is not None:

        output_values = [output_values[i] for i in self._prediction_outputs]

      if len(transformers) > 0:

    membership = multi_mol.membership

    deg_adjs = multi_mol.get_deg_adjacency_lists()[1:]

    args = [atom_features, degree_slice, membership] + deg_adjs

    print("HIIIIII")

    print("deg_adjs")

    print(deg_adjs)

    assert 0 == 1

    layer = layers.GraphConv(out_channels)

    result = layer(args)

    assert result.shape == (n_atoms, out_channels)