Converted more code to TF 2

    else:

      return None

  def make_iterator(self,

  def make_tf_dataset(self,

                      batch_size=100,

                      epochs=1,

                      deterministic=False,

                      pad_batches=False):

    """Create a tf.data.Iterator that iterates over the data in this Dataset.

    """Create a tf.data.Dataset that iterates over the data in this Dataset.

    The iterator's get_next() method returns a tuple of three tensors (X, y, w)

    which can be used to retrieve the features, labels, and weights respectively.

    Each value returned by the Dataset's iterator is a tuple of (X, y, w) for

    one batch.

    Parameters

    ----------

              tf.TensorShape([None] + list(y.shape)),

              tf.TensorShape([None] + list(w.shape)))

    # Create a Tensorflow Dataset and have it create an Iterator.

    # Create a Tensorflow Dataset.

    def gen_data():

      for epoch in range(epochs):

                                             pad_batches):

          yield (X, y, w)

    dataset = tf.data.Dataset.from_generator(gen_data, dtypes, shapes)

    return dataset.make_one_shot_iterator()

    return tf.data.Dataset.from_generator(gen_data, dtypes, shapes)

class NumpyDataset(Dataset):

    assert new_data.y.shape == (num_datapoints * num_datasets, num_tasks)

    assert len(new_data.tasks) == len(datasets[0].tasks)

  def test_make_iterator(self):

  def test_make_tf_dataset(self):

    """Test creating a Tensorflow Iterator from a Dataset."""

    X = np.random.random((100, 5))

    y = np.random.random((100, 1))

    dataset = dc.data.NumpyDataset(X, y)

    iterator = dataset.make_iterator(

    iterator = dataset.make_tf_dataset(

        batch_size=10, epochs=2, deterministic=True)

    next_element = iterator.get_next()

    with self.session() as sess:

      for i in range(20):

        batch_X, batch_y, batch_w = sess.run(next_element)

    for i, (batch_X, batch_y, batch_w) in enumerate(iterator):

      offset = (i % 10) * 10

      np.testing.assert_array_equal(X[offset:offset + 10, :], batch_X)

      np.testing.assert_array_equal(y[offset:offset + 10, :], batch_y)

      np.testing.assert_array_equal(np.ones((10, 1)), batch_w)

      finished = False

      try:

        sess.run(next_element)

      except tf.errors.OutOfRangeError:

        finished = True

    assert finished

    assert i == 19

if __name__ == "__main__":

    self.tensorboard = tensorboard

    self.tensorboard_log_frequency = tensorboard_log_frequency

    if self.tensorboard:

      if tf.executing_eagerly():

        raise ValueError(

            "Logging to TensorBoard is not currently supported in eager mode")

      self._summary_writer = tf.summary.FileWriter(self.model_dir)

      self._summary_writer = tf.summary.create_file_writer(self.model_dir)

    if output_types is None:

      self._prediction_outputs = None

      self._loss_outputs = None

        self.restore()

        restore = False

      inputs, labels, weights = self._prepare_batch(batch)

      self._current_summary = None

      # Execute the loss function, accumulating the gradients.

      for c in callbacks:

        c(self, current_step)

      if self.tensorboard and should_log:

        self._log_value_to_tensorboard(tag='loss', simple_value=batch_loss)

        self._summary_writer.reopen()

        self._summary_writer.add_summary(self._current_summary, current_step)

        self._summary_writer.close()

        with self._summary_writer.as_default():

          tf.summary.scalar('loss', batch_loss, current_step)

    # Report final results.

    if averaged_batches > 0:

    logger.info("TIMING: model fitting took %0.3f s" % (time2 - time1))

    return avg_loss

  def _log_value_to_tensorboard(self, **kwargs):

    """This can be called during fitting to log a value to Tensorboard.

    Any keyword arguments passed to this method are passed on to summary.value.add().

    """

    if self._current_summary is None:

      self._current_summary = tf.Summary()

    self._current_summary.value.add(**kwargs)

  def fit_on_batch(self, X, y, w, variables=None, loss=None, callbacks=[]):

    """Perform a single step of training.

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

  """Generate Gaussian nose."""

  def __init__(self, training_only=False, noise_epsilon=0.01, **kwargs):

  def __init__(self, training_only=False, noise_epsilon=1.0, **kwargs):

    """Create a CombineMeanStd layer.

    This layer should have two inputs with the same shape, and its output also has the

      if True, noise is only generated during training.  During prediction, the output

      is simply equal to the first input (that is, the mean of the distribution used

      during training).

    noise_epsilon: float

      The noise is scaled by this factor

    """

    super(CombineMeanStd, self).__init__(**kwargs)

    self.training_only = training_only

    noise_scale = tf.cast(training or not self.training_only, tf.float32)

    from tensorflow.python.ops import array_ops

    sample_noise = tf.random.normal(

        array_ops.shape(mean_parent), 0, 1, dtype=tf.float32)

        array_ops.shape(mean_parent), 0, self.noise_epsilon, dtype=tf.float32)

    return mean_parent + noise_scale * std_parent * sample_noise

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

  """Output a trainable value."""

  """Output a trainable value.

  Due to a quirk of Keras, you must pass an input value when invoking this layer.

  It doesn't matter what value you pass.  Keras assumes every layer that is not

  an Input will have at least one parent, and violating this assumption causes

  errors during evaluation.

  """

  def __init__(self, initial_value, **kwargs):

    """Construct a variable layer.

    self.num_outputs = len(inputs)

    # create subspaces

    subspaces = []

    original_cols = int(inputs[0].get_shape()[-1].value)

    original_cols = int(inputs[0].get_shape()[-1])

    subspace_size = int(original_cols / 2)

    for input_tensor in inputs:

      subspaces.append(tf.reshape(input_tensor[:, :subspace_size], [-1]))

    subspaces = []

    # creates subspaces the same way it was done in AlphaShare

    for input_tensor in inputs:

      subspace_size = int(input_tensor.get_shape()[-1].value / 2)

      subspace_size = int(input_tensor.get_shape()[-1] / 2)

      subspaces.append(input_tensor[:, :subspace_size])

      subspaces.append(input_tensor[:, subspace_size:])

      product = tf.matmul(tf.transpose(subspaces[0]), subspaces[1])

    Size of input layers must all be the same

    """

    subspaces = []

    original_cols = int(inputs[0].get_shape()[-1].value)

    original_cols = int(inputs[0].get_shape()[-1])

    for input_tensor in inputs:

      subspaces.append(tf.reshape(input_tensor, [-1]))

    n_betas = len(inputs)

    # Calculate pairwise distance

    d = tf.sqrt(

        tf.reduce_sum(tf.squared_difference(tensor1, tensor2), axis=3) + 1e-7)

        tf.reduce_sum(tf.math.squared_difference(tensor1, tensor2), axis=3) +

        1e-7)

    d = d * flags

    return d

  def build(self, input_shape):

    no_features = int(input_shape[0][-1])

    self.W = tf.Variable(

        tf.truncated_normal(

        tf.random.truncated_normal(

            [no_features, self.num_vertices],

            stddev=1.0 / np.sqrt(no_features)),

        name='weights',

    result_A = tf.reshape(result_A, (tf.shape(A)[0], tf.shape(A)[-1], -1))

    result_A = tf.matmul(factors, result_A, transpose_a=True)

    result_A = tf.reshape(result_A, (tf.shape(A)[0], self.num_vertices,

                                     A.get_shape()[2].value, self.num_vertices))

                                     A.get_shape()[2], self.num_vertices))

    return result, result_A

  def embedding_factors(self, V):

    no_features = V.get_shape()[-1].value

    no_features = V.get_shape()[-1]

    V_reshape = tf.reshape(V, (-1, no_features))

    s = tf.slice(tf.shape(V), [0], [len(V.get_shape()) - 1])

    s = tf.concat([s, tf.stack([self.num_vertices])], 0)

    no_features = int(input_shape[0][2])

    no_A = int(input_shape[1][2])

    self.W = tf.Variable(

        tf.truncated_normal(

        tf.random.truncated_normal(

            [no_features * no_A, self.num_filters],

            stddev=np.sqrt(1.0 / (no_features * (no_A + 1) * 1.0))),

        name='weights',

        dtype=tf.float32)

    self.W_I = tf.Variable(

        tf.truncated_normal(

        tf.random.truncated_normal(

            [no_features, self.num_filters],

            stddev=np.sqrt(1.0 / (no_features * (no_A + 1) * 1.0))),

        name='weights_I',

      V, A, mask = inputs

    else:

      V, A = inputs

    no_A = A.get_shape()[2].value

    no_features = V.get_shape()[2].value

    no_A = A.get_shape()[2]

    no_features = V.get_shape()[2]

    n = self.graphConvolution(V, A)

    A_shape = tf.shape(A)

    n = tf.reshape(n, [-1, A_shape[1], no_A * no_features])

        V, self.W_I) + self.b

  def graphConvolution(self, V, A):

    no_A = A.get_shape()[2].value

    no_features = V.get_shape()[2].value

    no_A = A.get_shape()[2]

    no_features = V.get_shape()[2]

    A_shape = tf.shape(A)

    A_reshape = tf.reshape(A, tf.stack([-1, A_shape[1] * no_A, A_shape[1]]))

    n = tf.matmul(A_reshape, V)

        tf.random.truncated_normal((1,), stddev=alpha_init_stddev))

    trainable_layers.append(alpha)

    prev_layer = Multiply()([prev_layer, alpha([])])

    prev_layer = Multiply()([prev_layer, alpha([prev_layer])])

    dense1 = Dense(

        layer_sizes[i - 1],

        kernel_initializer=tf.keras.initializers.TruncatedNormal(