def test_training_methods(self):
    # test fit, train_on_batch
    # on different input types: list, dict

    num_classes = (2, 3)
    num_samples = 100
    input_dim = 50

    x1 = np.ones((num_samples, input_dim))
    x2 = np.ones((num_samples, input_dim))
    y1 = np.zeros((num_samples, num_classes[0]))
    y2 = np.zeros((num_samples, num_classes[1]))

    model = MultiIOTestModel(num_classes=num_classes, use_bn=True)
    model.compile(
        loss='mse',
        optimizer='rmsprop',
        run_eagerly=testing_utils.should_run_eagerly())
    model.fit([x1, x2], [y1, y2], epochs=2, batch_size=32, verbose=0)
    model.fit({'input_1': x1, 'input_2': x2},
              {'output_1': y1, 'output_2': y2},
              epochs=2, batch_size=32)
    model.fit([x1, x2], [y1, y2], epochs=2, batch_size=32, verbose=0,
              validation_data=([x1, x2], [y1, y2]))

    model = MultiIOTestModel(num_classes=num_classes, use_bn=True)
    model.compile(
        loss='mse',
        optimizer='rmsprop',
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch([x1, x2], [y1, y2])
    model.train_on_batch({'input_1': x1, 'input_2': x2},
                         {'output_1': y1, 'output_2': y2})

  def test_merge_dot(self):
    i1 = keras.layers.Input(shape=(4,))
    i2 = keras.layers.Input(shape=(4,))
    o = keras.layers.dot([i1, i2], axes=1)
    self.assertListEqual(o.shape.as_list(), [None, 1])
    model = keras.models.Model([i1, i2], o)
    model.run_eagerly = testing_utils.should_run_eagerly()
    _ = keras.layers.Dot(axes=1).get_config()

    x1 = np.random.random((2, 4))
    x2 = np.random.random((2, 4))
    out = model.predict([x1, x2])
    self.assertEqual(out.shape, (2, 1))
    expected = np.zeros((2, 1))
    expected[0, 0] = np.dot(x1[0], x2[0])
    expected[1, 0] = np.dot(x1[1], x2[1])
    self.assertAllClose(out, expected, atol=1e-4)

    # Test with negative tuple of axes.
    o = keras.layers.dot([i1, i2], axes=(-1, -1))
    self.assertListEqual(o.shape.as_list(), [None, 1])
    model = keras.models.Model([i1, i2], o)
    model.run_eagerly = testing_utils.should_run_eagerly()
    out = model.predict([x1, x2])
    self.assertEqual(out.shape, (2, 1))
    self.assertAllClose(out, expected, atol=1e-4)

    # test compute_output_shape
    layer = keras.layers.Dot(axes=-1)
    self.assertEqual(layer.compute_output_shape([(4, 5), (4, 5)]), (4, 1))

  def test_multi_output_model_with_none_masking(self):
    def func(x):
      return [x * 0.2, x * 0.3]

    def output_shape(input_shape):
      return [input_shape, input_shape]

    i = keras.layers.Input(shape=(3, 2, 1))
    o = keras.layers.Lambda(function=func, output_shape=output_shape)(i)

    self.assertEqual(keras.backend.int_shape(o[0]), (None, 3, 2, 1))
    self.assertEqual(keras.backend.int_shape(o[1]), (None, 3, 2, 1))

    o = keras.layers.add(o)
    model = keras.Model(i, o)
    model.run_eagerly = testing_utils.should_run_eagerly()

    i2 = keras.layers.Input(shape=(3, 2, 1))
    o2 = model(i2)
    model2 = keras.Model(i2, o2)
    model2.run_eagerly = testing_utils.should_run_eagerly()

    x = np.random.random((4, 3, 2, 1))
    out = model2.predict(x)
    assert out.shape == (4, 3, 2, 1)
    self.assertAllClose(out, x * 0.2 + x * 0.3, atol=1e-4)

  def test_sequential_pop(self):
    num_hidden = 5
    input_dim = 3
    batch_size = 5
    num_classes = 2

    model = testing_utils.get_small_sequential_mlp(
        num_hidden, num_classes, input_dim)
    model.compile(loss='mse', optimizer=rmsprop.RMSPropOptimizer(1e-3),
                  run_eagerly=testing_utils.should_run_eagerly())
    x = np.random.random((batch_size, input_dim))
    y = np.random.random((batch_size, num_classes))
    model.fit(x, y, epochs=1)
    model.pop()
    self.assertEqual(len(model.layers), 1)
    self.assertEqual(model.output_shape, (None, num_hidden))
    model.compile(loss='mse', optimizer=rmsprop.RMSPropOptimizer(1e-3),
                  run_eagerly=testing_utils.should_run_eagerly())
    y = np.random.random((batch_size, num_hidden))
    model.fit(x, y, epochs=1)

    # Test popping single-layer model
    model = keras.models.Sequential()
    model.add(keras.layers.Dense(num_hidden, input_dim=input_dim))
    model.pop()
    self.assertEqual(model.layers, [])
    self.assertEqual(model.outputs, None)

    # Invalid use case
    model = keras.models.Sequential()
    with self.assertRaises(TypeError):
      model.pop()

  def test_nested_input_output_with_state(self):
    batch = 10
    t = 5
    i1, i2, i3 = 3, 4, 5
    o1, o2, o3 = 2, 3, 4

    cell = NestedCell(o1, o2, o3)
    rnn = keras.layers.RNN(cell, return_sequences=True, return_state=True)

    input_1 = keras.Input((t, i1))
    input_2 = keras.Input((t, i2, i3))

    output1, output2, s1, s2 = rnn((input_1, input_2))

    self.assertEqual(output1.shape.as_list(), [None, t, o1])
    self.assertEqual(output2.shape.as_list(), [None, t, o2, o3])
    self.assertEqual(s1.shape.as_list(), [None, o1])
    self.assertEqual(s2.shape.as_list(), [None, o2, o3])

    model = keras.models.Model([input_1, input_2], [output1, output2])
    model.compile(
        optimizer='rmsprop',
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(
        [np.zeros((batch, t, i1)),
         np.zeros((batch, t, i2, i3))],
        [np.zeros((batch, t, o1)),
         np.zeros((batch, t, o2, o3))])
    self.assertEqual(model.output_shape, [(None, t, o1), (None, t, o2, o3)])

    cell = NestedCell(o1, o2, o3, use_tuple=True)

    rnn = keras.layers.RNN(cell, return_sequences=True, return_state=True)

    input_1 = keras.Input((t, i1))
    input_2 = keras.Input((t, i2, i3))

    output1, output2, s1, s2 = rnn(NestedInput(t1=input_1, t2=input_2))

    self.assertEqual(output1.shape.as_list(), [None, t, o1])
    self.assertEqual(output2.shape.as_list(), [None, t, o2, o3])
    self.assertEqual(s1.shape.as_list(), [None, o1])
    self.assertEqual(s2.shape.as_list(), [None, o2, o3])

    model = keras.models.Model([input_1, input_2], [output1, output2])
    model.compile(
        optimizer='rmsprop',
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(
        [np.zeros((batch, t, i1)),
         np.zeros((batch, t, i2, i3))],
        [np.zeros((batch, t, o1)),
         np.zeros((batch, t, o2, o3))])
    self.assertEqual(model.output_shape, [(None, t, o1), (None, t, o2, o3)])

  def test_clone_functional_model(self, share_weights):
    if share_weights:
      clone_fn = functools.partial(
          keras.models._clone_functional_model, share_weights=True)
    else:
      clone_fn = keras.models.clone_model

    val_a = np.random.random((10, 4))
    val_b = np.random.random((10, 4))
    val_out = np.random.random((10, 4))

    input_a = keras.Input(shape=(4,))
    input_b = keras.Input(shape=(4,))
    dense_1 = keras.layers.Dense(4,)
    dense_2 = keras.layers.Dense(4,)

    x_a = dense_1(input_a)
    x_a = keras.layers.Dropout(0.5)(x_a)
    x_a = keras.layers.BatchNormalization()(x_a)
    x_b = dense_1(input_b)
    x_a = dense_2(x_a)
    outputs = keras.layers.add([x_a, x_b])
    model = keras.models.Model([input_a, input_b], outputs)

    # With placeholder creation
    new_model = clone_fn(model)
    self.assertEqual(len(new_model.get_updates_for(new_model.inputs)), 2)
    new_model.compile(
        testing_utils.get_v2_optimizer('rmsprop'), 'mse',
        run_eagerly=testing_utils.should_run_eagerly())
    new_model.train_on_batch([val_a, val_b], val_out)

    # On top of new tensors
    input_a = keras.Input(shape=(4,), name='a')
    input_b = keras.Input(shape=(4,), name='b')
    new_model = keras.models.clone_model(
        model, input_tensors=[input_a, input_b])
    self.assertEqual(len(new_model.get_updates_for(new_model.inputs)), 2)
    new_model.compile(
        testing_utils.get_v2_optimizer('rmsprop'), 'mse',
        run_eagerly=testing_utils.should_run_eagerly())
    new_model.train_on_batch([val_a, val_b], val_out)

    # On top of new, non-Keras tensors
    if not context.executing_eagerly():
      # TODO(b/121277734):Skip Eager contexts, as Input() layers raise an error
      # saying they should not be used with EagerTensors
      input_a = keras.backend.variable(val_a)
      input_b = keras.backend.variable(val_b)
      new_model = clone_fn(model, input_tensors=[input_a, input_b])
      self.assertEqual(len(new_model.get_updates_for(new_model.inputs)), 2)
      new_model.compile(
          testing_utils.get_v2_optimizer('rmsprop'), 'mse',
          run_eagerly=testing_utils.should_run_eagerly())
      new_model.train_on_batch(None, val_out)

  def test_minimal_rnn_cell_abstract_rnn_cell(self):

    class MinimalRNNCell(keras.layers.AbstractRNNCell):

      def __init__(self, units, **kwargs):
        self.units = units
        super(MinimalRNNCell, self).__init__(**kwargs)

      @property
      def state_size(self):
        return self.units

      def build(self, input_shape):
        self.kernel = self.add_weight(shape=(input_shape[-1], self.units),
                                      initializer='uniform',
                                      name='kernel')
        self.recurrent_kernel = self.add_weight(
            shape=(self.units, self.units),
            initializer='uniform',
            name='recurrent_kernel')
        self.built = True

      def call(self, inputs, states):
        prev_output = states[0]
        h = keras.backend.dot(inputs, self.kernel)
        output = h + keras.backend.dot(prev_output, self.recurrent_kernel)
        return output, output

      @property
      def output_size(self):
        return self.units

    cell = MinimalRNNCell(32)
    x = keras.Input((None, 5))
    layer = keras.layers.RNN(cell)
    y = layer(x)
    model = keras.models.Model(x, y)
    model.compile(
        optimizer="rmsprop",
        loss="mse",
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(np.zeros((6, 5, 5)), np.zeros((6, 32)))

    # Test stacking.
    cells = [MinimalRNNCell(8),
             MinimalRNNCell(16),
             MinimalRNNCell(32)]
    layer = keras.layers.RNN(cells)
    y = layer(x)
    model = keras.models.Model(x, y)
    model.compile(optimizer='rmsprop',
                  loss='mse',
                  run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(np.zeros((6, 5, 5)), np.zeros((6, 32)))

  def test_nested_input_output(self):
    batch = 10
    t = 5
    i1, i2, i3 = 3, 4, 5
    o1, o2, o3 = 2, 3, 4

    cell = NestedCell(o1, o2, o3)
    rnn = keras.layers.RNN(cell)

    input_1 = keras.Input((t, i1))
    input_2 = keras.Input((t, i2, i3))

    outputs = rnn((input_1, input_2))

    self.assertEqual(len(outputs), 2)
    self.assertEqual(outputs[0].shape.as_list(), [None, o1])
    self.assertEqual(outputs[1].shape.as_list(), [None, o2, o3])

    model = keras.models.Model((input_1, input_2), outputs)
    model.compile(
        optimizer='rmsprop',
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(
        [np.zeros((batch, t, i1)), np.zeros((batch, t, i2, i3))],
        [np.zeros((batch, o1)), np.zeros((batch, o2, o3))])
    self.assertEqual(model.output_shape, [(None, o1), (None, o2, o3)])

    cell = NestedCell(o1, o2, o3, use_tuple=True)

    rnn = keras.layers.RNN(cell)

    input_1 = keras.Input((t, i1))
    input_2 = keras.Input((t, i2, i3))

    outputs = rnn(NestedInput(t1=input_1, t2=input_2))

    self.assertEqual(len(outputs), 2)
    self.assertEqual(outputs[0].shape.as_list(), [None, o1])
    self.assertEqual(outputs[1].shape.as_list(), [None, o2, o3])

    model = keras.models.Model([input_1, input_2], outputs)
    model.compile(
        optimizer='rmsprop',
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(
        [np.zeros((batch, t, i1)),
         np.zeros((batch, t, i2, i3))],
        [np.zeros((batch, o1)), np.zeros((batch, o2, o3))])
    self.assertEqual(model.output_shape, [(None, o1), (None, o2, o3)])

  def test_builtin_rnn_cell_serialization(self):
    for cell_class in [keras.layers.SimpleRNNCell,
                       keras.layers.GRUCell,
                       keras.layers.LSTMCell]:
      # Test basic case.
      x = keras.Input((None, 5))
      cell = cell_class(32)
      layer = keras.layers.RNN(cell)
      y = layer(x)
      model = keras.models.Model(x, y)
      model.compile(
          optimizer='rmsprop',
          loss='mse',
          run_eagerly=testing_utils.should_run_eagerly())

      # Test basic case serialization.
      x_np = np.random.random((6, 5, 5))
      y_np = model.predict(x_np)
      weights = model.get_weights()
      config = layer.get_config()
      layer = keras.layers.RNN.from_config(config)
      y = layer(x)
      model = keras.models.Model(x, y)
      model.set_weights(weights)
      y_np_2 = model.predict(x_np)
      self.assertAllClose(y_np, y_np_2, atol=1e-4)

      # Test stacking.
      cells = [cell_class(8),
               cell_class(12),
               cell_class(32)]
      layer = keras.layers.RNN(cells)
      y = layer(x)
      model = keras.models.Model(x, y)
      model.compile(
          optimizer='rmsprop',
          loss='mse',
          run_eagerly=testing_utils.should_run_eagerly())

      # Test stacked RNN serialization.
      x_np = np.random.random((6, 5, 5))
      y_np = model.predict(x_np)
      weights = model.get_weights()
      config = layer.get_config()
      layer = keras.layers.RNN.from_config(config)
      y = layer(x)
      model = keras.models.Model(x, y)
      model.set_weights(weights)
      y_np_2 = model.predict(x_np)
      self.assertAllClose(y_np, y_np_2, atol=1e-4)

  def test_high_dimension_RNN(self):
    # Basic test case.
    unit_a = 10
    unit_b = 20
    input_a = 5
    input_b = 10
    batch = 32
    time_step = 4

    cell = Minimal2DRNNCell(unit_a, unit_b)
    x = keras.Input((None, input_a, input_b))
    layer = keras.layers.RNN(cell)
    y = layer(x)

    self.assertEqual(cell.state_size.as_list(), [unit_a, unit_b])

    if not context.executing_eagerly():
      init_state = layer.get_initial_state(x)
      self.assertEqual(len(init_state), 1)
      self.assertEqual(init_state[0].get_shape().as_list(),
                       [None, unit_a, unit_b])

    model = keras.models.Model(x, y)
    model.compile(
        optimizer='rmsprop',
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(
        np.zeros((batch, time_step, input_a, input_b)),
        np.zeros((batch, unit_a, unit_b)))
    self.assertEqual(model.output_shape, (None, unit_a, unit_b))

    # Test stacking.
    cells = [
        Minimal2DRNNCell(unit_a, unit_b),
        Minimal2DRNNCell(unit_a * 2, unit_b * 2),
        Minimal2DRNNCell(unit_a * 4, unit_b * 4)
    ]
    layer = keras.layers.RNN(cells)
    y = layer(x)
    model = keras.models.Model(x, y)
    model.compile(
        optimizer='rmsprop',
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(
        np.zeros((batch, time_step, input_a, input_b)),
        np.zeros((batch, unit_a * 4, unit_b * 4)))
    self.assertEqual(model.output_shape, (None, unit_a * 4, unit_b * 4))

  def test_metrics_correctness_with_iterator(self):
    layers = [
        keras.layers.Dense(8, activation='relu', input_dim=4,
                           kernel_initializer='ones'),
        keras.layers.Dense(1, activation='sigmoid', kernel_initializer='ones')
    ]

    model = testing_utils.get_model_from_layers(layers, (4,))

    model.compile(
        loss='binary_crossentropy',
        metrics=['accuracy', metrics_module.BinaryAccuracy()],
        optimizer='rmsprop',
        run_eagerly=testing_utils.should_run_eagerly())

    np.random.seed(123)
    x = np.random.randint(10, size=(100, 4)).astype(np.float32)
    y = np.random.randint(2, size=(100, 1)).astype(np.float32)
    dataset = dataset_ops.Dataset.from_tensor_slices((x, y))
    dataset = dataset.batch(10)
    iterator = dataset_ops.make_one_shot_iterator(dataset)
    outs = model.evaluate(iterator, steps=10)
    self.assertEqual(np.around(outs[1], decimals=1), 0.5)
    self.assertEqual(np.around(outs[2], decimals=1), 0.5)

    y = np.zeros((100, 1), dtype=np.float32)
    dataset = dataset_ops.Dataset.from_tensor_slices((x, y))
    dataset = dataset.repeat(100)
    dataset = dataset.batch(10)
    iterator = dataset_ops.make_one_shot_iterator(dataset)
    outs = model.evaluate(iterator, steps=10)
    self.assertEqual(outs[1], 0.)
    self.assertEqual(outs[2], 0.)

  def test_finite_dataset_unknown_cardinality_out_of_data(self):
    model = testing_utils.get_small_mlp(1, 4, input_dim=3)
    model.compile('rmsprop', 'mse',
                  run_eagerly=testing_utils.should_run_eagerly())

    inputs = np.zeros((100, 3), dtype=np.float32)
    targets = np.random.randint(0, 4, size=100, dtype=np.int32)
    dataset = dataset_ops.Dataset.from_tensor_slices((inputs, targets))
    dataset = dataset.filter(lambda x, y: True).batch(10)
    self.assertEqual(
        keras.backend.get_value(cardinality.cardinality(dataset)),
        cardinality.UNKNOWN)

    batch_counter = BatchCounterCallback()
    with test.mock.patch.object(logging, 'warning') as mock_log:
      # steps_per_epoch (200) is greater than the dataset size (100). As this is
      # unexpected, training will stop and not make it to the second epoch.
      history = model.fit(
          dataset,
          epochs=2,
          verbose=1,
          callbacks=[batch_counter],
          steps_per_epoch=200)
      self.assertIn(
          'Your dataset ran out of data; interrupting training. '
          'Make sure that your dataset can generate at least '
          '`steps_per_epoch * epochs` batches (in this case, 400 batches). '
          'You may need to use the repeat() function when '
          'building your dataset.', str(mock_log.call_args))

    self.assertLen(history.history['loss'], 1)
    self.assertEqual(batch_counter.batch_count, 10)
    model.evaluate(dataset)
    out = model.predict(dataset)
    self.assertEqual(out.shape[0], 100)

  def test_vector_classification(self):
    np.random.seed(1337)
    (x_train, y_train), _ = testing_utils.get_test_data(
        train_samples=100,
        test_samples=0,
        input_shape=(10,),
        num_classes=2)
    y_train = keras.utils.to_categorical(y_train)

    model = testing_utils.get_model_from_layers(
        [keras.layers.Dense(16, activation='relu'),
         keras.layers.Dropout(0.1),
         keras.layers.Dense(y_train.shape[-1], activation='softmax')],
        input_shape=x_train.shape[1:])
    model.compile(
        loss='categorical_crossentropy',
        optimizer=keras.optimizer_v2.adam.Adam(0.005),
        metrics=['acc'],
        run_eagerly=testing_utils.should_run_eagerly())
    history = model.fit(x_train, y_train, epochs=10, batch_size=10,
                        validation_data=(x_train, y_train),
                        verbose=2)
    self.assertGreater(history.history['val_acc'][-1], 0.7)
    _, val_acc = model.evaluate(x_train, y_train)
    self.assertAlmostEqual(history.history['val_acc'][-1], val_acc)
    predictions = model.predict(x_train)
    self.assertEqual(predictions.shape, (x_train.shape[0], 2))

  def test_merge_subtract(self):
    i1 = keras.layers.Input(shape=(4, 5))
    i2 = keras.layers.Input(shape=(4, 5))
    i3 = keras.layers.Input(shape=(4, 5))

    subtract_layer = keras.layers.Subtract()
    o = subtract_layer([i1, i2])
    self.assertListEqual(o.shape.as_list(), [None, 4, 5])
    model = keras.models.Model([i1, i2], o)
    model.run_eagerly = testing_utils.should_run_eagerly()

    x1 = np.random.random((2, 4, 5))
    x2 = np.random.random((2, 4, 5))
    out = model.predict([x1, x2])
    self.assertEqual(out.shape, (2, 4, 5))
    self.assertAllClose(out, x1 - x2, atol=1e-4)

    self.assertEqual(subtract_layer.compute_mask([i1, i2], [None, None]), None)
    self.assertTrue(
        np.all(
            K.eval(
                subtract_layer.compute_mask(
                    [i1, i2], [K.variable(x1), K.variable(x2)]))))

    with self.assertRaisesRegexp(ValueError, "`mask` should be a list."):
      subtract_layer.compute_mask([i1, i2], x1)
    with self.assertRaisesRegexp(ValueError, "`inputs` should be a list."):
      subtract_layer.compute_mask(i1, [None, None])
    with self.assertRaisesRegexp(ValueError,
                                 "layer should be called on exactly 2 inputs"):
      subtract_layer([i1, i2, i3])
    with self.assertRaisesRegexp(ValueError,
                                 "layer should be called on exactly 2 inputs"):
      subtract_layer([i1])

  def test_sequential_deferred_build_serialization(self):
    num_hidden = 5
    input_dim = 3
    batch_size = 5
    num_classes = 2

    model = testing_utils.get_small_sequential_mlp(num_hidden, num_classes)
    model.compile(
        loss='mse',
        optimizer='rmsprop',
        metrics=[keras.metrics.CategoricalAccuracy()],
        run_eagerly=testing_utils.should_run_eagerly())
    self.assertFalse(model.built)

    x = np.random.random((batch_size, input_dim))
    y = np.random.random((batch_size, num_classes))
    model.train_on_batch(x, y)
    self.assertTrue(model.built)

    config = model.get_config()
    self.assertIn('build_input_shape', config)

    new_model = keras.models.Sequential.from_config(config)
    self.assertEqual(len(new_model.layers), 2)
    self.assertEqual(len(new_model.weights), 4)

  def test_updates(self):
    # test that updates get run during training
    num_samples = 100
    input_dim = 50

    class BNNet(keras.Model):

      def __init__(self):
        super(BNNet, self).__init__()
        self.bn = keras.layers.BatchNormalization(beta_initializer='ones',
                                                  gamma_initializer='ones')

      def call(self, inputs):
        return self.bn(inputs)

    x = np.ones((num_samples, input_dim))
    y = np.ones((num_samples, input_dim))

    model = BNNet()
    model.compile(
        loss='mse',
        optimizer='rmsprop',
        run_eagerly=testing_utils.should_run_eagerly())
    y_ref = model.predict(x)

    model.train_on_batch(x, y)
    y_new = model.predict(x)
    self.assertGreater(np.sum(np.abs(y_ref - y_new)), 0.1)

  def test_inconsistent_output_state_size(self):
    batch = 32
    time_step = 4
    state_size = 5
    input_size = 6
    cell = PlusOneRNNCell(state_size)
    x = keras.Input((None, input_size))
    layer = keras.layers.RNN(cell)
    y = layer(x)

    self.assertEqual(cell.state_size, state_size)
    if not context.executing_eagerly():
      init_state = layer.get_initial_state(x)
      self.assertEqual(len(init_state), 1)
      self.assertEqual(init_state[0].get_shape().as_list(),
                       [None, state_size])

    model = keras.models.Model(x, y)
    model.compile(
        optimizer='rmsprop',
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly())
    model.train_on_batch(
        np.zeros((batch, time_step, input_size)),
        np.zeros((batch, input_size)))
    self.assertEqual(model.output_shape, (None, input_size))

  def test_zero_output_for_masking(self):

    for unroll in [True, False]:
      cell = keras.layers.SimpleRNNCell(5)
      x = keras.Input((5, 5))
      mask = keras.layers.Masking()
      layer = keras.layers.RNN(
          cell, return_sequences=True, zero_output_for_mask=True, unroll=unroll)
      masked_input = mask(x)
      y = layer(masked_input)
      model = keras.models.Model(x, y)
      model.compile(
          optimizer='rmsprop',
          loss='mse',
          run_eagerly=testing_utils.should_run_eagerly())

      np_x = np.ones((6, 5, 5))
      result_1 = model.predict(np_x)

      # set the time 4 and 5 for last record to be zero (masked).
      np_x[5, 3:] = 0
      result_2 = model.predict(np_x)

      # expect the result_2 has same output, except the time 4,5 for last
      # record.
      result_1[5, 3:] = 0
      self.assertAllClose(result_1, result_2)

  def test_sequential_as_downstream_of_masking_layer(self):
    inputs = keras.layers.Input(shape=(3, 4))
    x = keras.layers.Masking(mask_value=0., input_shape=(3, 4))(inputs)

    s = keras.Sequential()
    s.add(keras.layers.Dense(5, input_shape=(4,)))

    x = keras.layers.wrappers.TimeDistributed(s)(x)
    model = keras.Model(inputs=inputs, outputs=x)
    model.compile(
        optimizer='rmsprop',
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly())

    model_input = np.random.randint(
        low=1, high=5, size=(10, 3, 4)).astype('float32')
    for i in range(4):
      model_input[i, i:, :] = 0.
    model.fit(model_input,
              np.random.random((10, 3, 5)), epochs=1, batch_size=6)

    if not context.executing_eagerly():
      # Note: this doesn't work in eager due to DeferredTensor/ops compatibility
      # issue.
      mask_outputs = [model.layers[1].compute_mask(model.layers[1].input)]
      mask_outputs += [model.layers[2].compute_mask(
          model.layers[2].input, mask_outputs[-1])]
      func = keras.backend.function([model.input], mask_outputs)
      mask_outputs_val = func([model_input])
      self.assertAllClose(mask_outputs_val[0], np.any(model_input, axis=-1))
      self.assertAllClose(mask_outputs_val[1], np.any(model_input, axis=-1))

  def test_merge_concatenate(self):
    i1 = keras.layers.Input(shape=(4, 5))
    i2 = keras.layers.Input(shape=(4, 5))
    concat_layer = keras.layers.Concatenate(axis=1)
    o = concat_layer([i1, i2])
    self.assertListEqual(o.shape.as_list(), [None, 8, 5])
    model = keras.models.Model([i1, i2], o)
    model.run_eagerly = testing_utils.should_run_eagerly()

    x1 = np.random.random((2, 4, 5))
    x2 = np.random.random((2, 4, 5))
    out = model.predict([x1, x2])
    self.assertEqual(out.shape, (2, 8, 5))
    self.assertAllClose(out, np.concatenate([x1, x2], axis=1), atol=1e-4)

    self.assertEqual(concat_layer.compute_mask([i1, i2], [None, None]), None)
    self.assertTrue(
        np.all(
            K.eval(
                concat_layer.compute_mask(
                    [i1, i2], [K.variable(x1), K.variable(x2)]))))

    with self.assertRaisesRegexp(ValueError, "`mask` should be a list."):
      concat_layer.compute_mask([i1, i2], x1)
    with self.assertRaisesRegexp(ValueError, "`inputs` should be a list."):
      concat_layer.compute_mask(i1, [None, None])
    with self.assertRaisesRegexp(ValueError, "should have the same length"):
      concat_layer.compute_mask([i1, i2], [None])
    with self.assertRaisesRegexp(ValueError,
                                 "layer should be called on a list of inputs"):
      concat_layer(i1)