Converted graph layers from TensorGraph to Keras

    return out_tensor

class Highway(Layer):

class Highway(KerasLayer):

  """ Create a highway layer. y = H(x) * T(x) + x * (1 - T(x))

  H(x) = activation_fn(matmul(W_H, x) + b_H) is the non-linear transformed output

  T(x) = sigmoid(matmul(W_T, x) + b_T) is the transform gate

    except:

      pass

  def _build_layers(self, out_channels):

    if self.biases_initializer is None:

      biases_initializer = None

    else:

      biases_initializer = self.biases_initializer()

    dense_H = tf.keras.layers.Dense(

        out_channels,

        activation=self.activation_fn,

        bias_initializer=biases_initializer,

        kernel_initializer=self.weights_initializer())

    dense_T = tf.keras.layers.Dense(

        out_channels,

        activation=tf.nn.sigmoid,

        bias_initializer=tf.constant_initializer(-1),

        kernel_initializer=self.weights_initializer())

    return (dense_H, dense_T)

  def create_tensor(self, in_layers=None, set_tensors=True, **kwargs):

    inputs = self._get_input_tensors(in_layers)

    parent = inputs[0]

    out_channels = parent.get_shape().as_list()[1]

    if tf.executing_eagerly():

      if not self._built:

        self._layers = self._build_layers(out_channels)

        self._non_pickle_fields.append('_layers')

      layers = self._layers

    else:

      layers = self._build_layers(out_channels)

    dense_H = layers[0](parent)

    dense_T = layers[1](parent)

    out_tensor = tf.multiply(dense_H, dense_T) + tf.multiply(

        parent, 1 - dense_T)

    if set_tensors:

      self.out_tensor = out_tensor

      self.trainable_variables = layers[0].trainable_variables + layers[1].trainable_variables

    if tf.executing_eagerly() and not self._built:

      self._built = True

      self.trainable_variables = layers[0].trainable_variables + layers[1].trainable_variables

    return out_tensor

  def _build_layer(self):

    return deepchem.models.layers.Highway(

        self.activation_fn, self.biases_initializer, self.weights_initializer)

class Flatten(Layer):

from deepchem.models.tensorgraph.graph_layers import DTNNEmbedding

from deepchem.models.tensorgraph.graph_layers import DTNNExtract

from deepchem.models.tensorgraph.graph_layers import WeaveGather

from deepchem.models.tensorgraph.graph_layers import GatedRecurrentUnit

from deepchem.models.tensorgraph.graph_layers import EdgeNetwork

from deepchem.models.tensorgraph.graph_layers import DTNNGather

from deepchem.models.tensorgraph.graph_layers import DAGGather

from deepchem.models.layers import GatedRecurrentUnit

from deepchem.models.layers import EdgeNetwork

class TestLayers(test_util.TensorFlowTestCase):

      # Layer is wrapper around embedding lookup, tested that then

      sess.run(tf.global_variables_initializer())

      out_tensor = dtnn_embedding.out_tensor.eval()

      embedding_val = dtnn_embedding.embedding_list.eval()

      embedding_val = dtnn_embedding.trainable_variables[0].eval()

      expected_output = embedding_val[test_tensor_input]

      self.assertAllClose(out_tensor, expected_output)

      self.assertAllClose(out_tensor.shape,

      messages_tf = tf.convert_to_tensor(messages_np, dtype=tf.float32)

      inputs_tf = tf.convert_to_tensor(inputs_np, dtype=tf.float32)

      gru = GatedRecurrentUnit(n_hidden=n_hidden, init=init_method)

      h = gru.forward(inputs_tf, messages_tf)

      h = gru([inputs_tf, messages_tf])

      sess.run(tf.global_variables_initializer())

      h = h.eval()

          atom_features_np, dtype=tf.float32)

      atom_to_pair_tf = tf.convert_to_tensor(atom_to_pair_np, dtype=tf.int32)

      edge_network = EdgeNetwork(

          pair_features=pair_features_tf,

          n_pair_features=n_pair_features,

          n_hidden=n_hidden,

          init=init_method)

          n_pair_features=n_pair_features, n_hidden=n_hidden, init=init_method)

      edge_network.build([])

      sess.run(tf.global_variables_initializer())

      output = edge_network.forward(atom_features_tf, atom_to_pair_tf).eval()

      output = edge_network(

          [pair_features_tf, atom_features_tf, atom_to_pair_tf]).eval()

      self.assertAllClose(output, expected_output)

      self.assertEqual(output.shape, expected_output.shape)