Test fixes

          dropout = np.array(1.0)

        yield ([X_b, dropout], [y_b], [w_b])

  def predict_on_generator(self, generator, transformers=[], outputs=None):

  def predict_on_generator(self,

                           generator,

                           transformers=[],

                           outputs=None,

                           output_types=None):

    def transform_generator():

      for inputs, labels, weights in generator:

        yield ([X_t] + inputs[1:], labels, weights)

    return super(MultitaskFitTransformRegressor, self).predict_on_generator(

        transform_generator(), transformers, outputs)

        transform_generator(), transformers, outputs, output_types)

              -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

            np.array(multiConvMol.membership),

            n_samples

            multiConvMol.get_atom_features(), multiConvMol.deg_slice,

            np.array(multiConvMol.membership), n_samples

        ]

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

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

    def apply_gradient_for_batch(inputs, labels, weights, loss):

      with tf.GradientTape() as tape:

        outputs = self.model(inputs, training=True)

        #outputs = self.model(inputs)

        if isinstance(outputs, tf.Tensor):

          outputs = [outputs]

        if self._loss_outputs is not None:

      If True, it sets the training flag so that dropout will be enabled, and

      returns the values of the uncertainty outputs.

    other_output_types: list, optional

      Provides a list of other outputs to predict from model.

      Each such output should have a unique output_type so it

      can be retrieved from the model.

      Provides a list of other output_types (strings) to predict from model.

    Returns:

      a NumPy array of the model produces a single output, or a list of arrays

      if it produces multiple outputs

    """

    results = None

    variances = None

    if outputs and other_output_type:

    if (outputs is not None) and (other_output_types is not None):

      raise ValueError(

          'This model cannot compute outputs and other output_types simultaneously. Please invoke one at a time.'

      )

    if uncertainty and other_output_types:

    if uncertainty and (other_output_types is not None):

      raise ValueError(

          'This model cannot compute uncertainties and other output types simultaneously. Please invoke one at a time.'

      )

    """

    generator = self.default_generator(

        dataset, mode='predict', pad_batches=False)

    return self.predict_on_generator(generator, transformers, outputs,

                                     output_types)

    return self.predict_on_generator(

        generator,

        transformers=transformers,

        outputs=outputs,

        output_types=output_types)

  def predict_embedding(self, dataset):

    """

    return tf.math.segment_sum(output, atom_membership)

def _DAGgraph_step(

    batch_inputs,

    W_list,

    b_list,

    activation_fn,

    #dropout,

    dropouts,

    #dropout_switch):

def _DAGgraph_step(batch_inputs, W_list, b_list, activation_fn, dropouts,

                   training):

  outputs = batch_inputs

  #for idw, W in enumerate(W_list):

  for idw, (dropout, W) in enumerate(zip(dropouts, W_list)):

    outputs = tf.nn.bias_add(tf.matmul(outputs, W), b_list[idw])

    outputs = activation_fn(outputs)

    if dropout is not None:

      #outputs = tf.nn.dropout(outputs, rate=dropout * dropout_switch)

      outputs = dropout(outputs, training=training)

  return outputs

    calculation_masks = inputs[3]

    n_atoms = tf.squeeze(inputs[4])

    #dropout_switch = tf.squeeze(inputs[5])

    graph_features = tf.zeros((self.max_atoms * self.batch_size,

                               self.max_atoms + 1, self.n_graph_feat))

      # extracting graph features for parents of the target atoms, then flatten

      # shape: (batch_size*max_atoms) * [(max_atoms-1)*n_graph_features]

      batch_graph_features = tf.reshape(

          #tf.gather_nd(self.graph_features, index),

          tf.gather_nd(graph_features, index),

          [-1, (self.max_atoms - 1) * self.n_graph_feat])

    """

    atom_features = inputs[0]

    membership = inputs[1]

    #dropout_switch = tf.squeeze(inputs[2])

    # Extract atom_features

    graph_features = tf.math.segment_sum(atom_features, membership)

    # sum all graph outputs

    #return _DAGgraph_step(graph_features, self.W_list, self.b_list,

    #                      self.activation_fn, self.dropout, dropout_switch)

    return _DAGgraph_step(

        graph_features,

        self.W_list,

        self.b_list,

        self.activation_fn,

        #self.dropout,

        self.dropouts,

        #dropout_switch)

        training)

    return _DAGgraph_step(graph_features, self.W_list, self.b_list,

                          self.activation_fn, self.dropouts, training)

class MessagePassing(tf.keras.layers.Layer):

        layer_sizes=layer_sizes)

    atom_features = np.random.rand(batch_size, n_atom_feat)

    membership = np.sort(np.random.randint(0, batch_size, size=(batch_size)))

    #dropout_switch = False

    #outputs = layer([atom_features, membership, dropout_switch])

    outputs = layer([atom_features, membership])