text: # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Tests for tf.keras models using tf.distribute.Strategy."""
  
  import tensorflow.compat.v2 as tf
  
  import os
  
  from absl.testing import parameterized
  import numpy as np
  
  import keras
  from tensorflow.python.distribute.cluster_resolver import SimpleClusterResolver
  from keras import backend
  from keras import testing_utils
  from keras.distribute import distributed_training_utils
  from keras.distribute import distributed_training_utils_v1
  from keras.distribute import multi_worker_testing_utils
  from keras.distribute import optimizer_combinations
  from keras.distribute.strategy_combinations import all_strategies
  from keras.distribute.strategy_combinations import multi_worker_mirrored_strategies
  from keras.distribute.strategy_combinations import strategies_minus_default_minus_tpu
  from keras.distribute.strategy_combinations import strategies_minus_tpu
  from keras.distribute.strategy_combinations import tpu_strategies
  from keras.engine import base_layer_utils
  from keras.mixed_precision import policy
  from keras.optimizer_v2 import gradient_descent as gradient_descent_keras
  from keras.utils import losses_utils
  from keras.utils import np_utils
  
  _RANDOM_SEED = 1337
  _TRAIN_SIZE = 200
  _INPUT_SIZE = (10,)
  _NUM_CLASS = 2
  
  # Note: Please make sure the tests in this file are also covered in
  # keras_backward_compat_test for features that are supported with both APIs.
  
  # TODO(anjalisridhar): Add a decorator that will allow us to run these tests as
  # part of the tf.keras unit tests suite.
  
  
  def simple_sequential_model():
    model = keras.models.Sequential()
    model.add(keras.layers.Dense(16, activation='relu', input_shape=_INPUT_SIZE))
    model.add(keras.layers.Dropout(0.1))
    model.add(keras.layers.Dense(_NUM_CLASS, activation='softmax'))
    return model
  
  
  def simple_subclassed_model(num_labels=_NUM_CLASS):
  
    class _SimpleMLP(keras.Model):
  
      def __init__(self, num_labels):
        super(_SimpleMLP, self).__init__()
        self.dense = keras.layers.Dense(num_labels)
  
      def call(self, inputs):
        return self.dense(inputs)
  
    return _SimpleMLP(num_labels)
  
  
  def simple_multi_inputs_multi_outputs_model():
    input_a = keras.layers.Input(shape=(16,), name='input_a')
    input_b = keras.layers.Input(shape=(16,), name='input_b')
  
    merged = keras.layers.concatenate([input_a, input_b], name='merge')
    output_c = keras.layers.Dense(3, activation='softmax', name='dense_2')(merged)
    output_d = keras.layers.Dense(2, activation='softmax', name='dense_3')(merged)
    model = keras.models.Model(
        inputs=[input_a, input_b], outputs=[output_c, output_d])
    return model
  
  
  def get_multi_inputs_multi_outputs_data():
    (a_train, c_train), (a_test, c_test) = testing_utils.get_test_data(
        train_samples=_TRAIN_SIZE,
        test_samples=50,
        input_shape=(16,),
        num_classes=3,
        random_seed=_RANDOM_SEED)
    (b_train, d_train), (b_test, d_test) = testing_utils.get_test_data(
        train_samples=_TRAIN_SIZE,
        test_samples=50,
        input_shape=(16,),
        num_classes=2,
        random_seed=_RANDOM_SEED)
    (m_train, _), (m_test, _) = testing_utils.get_test_data(
        train_samples=_TRAIN_SIZE,
        test_samples=50,
        input_shape=(8,),
        num_classes=2,
        random_seed=_RANDOM_SEED)
  
    c_train = np_utils.to_categorical(c_train)
    c_test = np_utils.to_categorical(c_test)
    d_train = np_utils.to_categorical(d_train)
    d_test = np_utils.to_categorical(d_test)
  
    train_data = {
        'input_a': a_train,
        'input_b': b_train,
        'input_m': m_train,
        'output_c': c_train,
        'output_d': d_train
    }
    test_data = {
        'input_a': a_test,
        'input_b': b_test,
        'input_m': m_test,
        'output_c': c_test,
        'output_d': d_test
    }
  
    return (train_data, test_data)
  
  
  def batch_wrapper(dataset, batch_size, distribution, repeat=None):
    if repeat:
      dataset = dataset.repeat(repeat)
    # TPUs currently require fully defined input shapes, drop_remainder ensures
    # the input will have fully defined shapes.
    if backend.is_tpu_strategy(distribution):
      return dataset.batch(batch_size, drop_remainder=True)
    else:
      return dataset.batch(batch_size)
  
  
  def get_model():
    x = keras.layers.Input(shape=(3,), name='input')
    y = keras.layers.Dense(4, name='dense')(x)
    model = keras.Model(x, y)
    return model
  
  
  def get_sample_weights_model():
    x = keras.layers.Input(shape=(1,), name='input')
    y = keras.layers.Dense(
        1, kernel_initializer='ones', bias_initializer='zeros', name='dense')(
            x)
    model = keras.Model(x, y)
    return model
  
  
  def get_dataset(distribution):
    inputs = np.zeros((10, 3), dtype=np.float32)
    targets = np.zeros((10, 4), dtype=np.float32)
    dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
    dataset = dataset.repeat(100)
    dataset = batch_wrapper(dataset, 10, distribution)
    return dataset
  
  
  def get_predict_dataset(distribution):
    inputs = np.zeros((10, 3), dtype=np.float32)
    dataset = tf.data.Dataset.from_tensor_slices(inputs)
    dataset = dataset.repeat(100)
    dataset = batch_wrapper(dataset, 10, distribution)
    return dataset
  
  
  def convert_numpy_to_dataset_with_unknown_cardinality(inputs, targets=None):
    if targets is not None:
      input_slices = (inputs, targets)
      dummy_op = (lambda inp, target: True)
    else:
      input_slices = inputs
      dummy_op = (lambda inp: True)
  
    original_dataset = (tf.data.Dataset.from_tensor_slices(input_slices))
    ds_with_unknown_cardinality = (
        original_dataset.filter(dummy_op).batch(10, drop_remainder=True))
    return ds_with_unknown_cardinality
  
  
  def multi_input_output_model():
    a = keras.layers.Input(shape=(3,), name='input_a')
    b = keras.layers.Input(shape=(5,), name='input_b')
    # TODO(anjalisridhar): Change the output dimension of the second Dense layer
    # once the iterator output validation issue has been fixed.
    dense_1 = keras.layers.Dense(7, name='dense_1')
    dense_2 = keras.layers.Dense(7, name='dense_2')
    c = dense_1(a)
    d = dense_2(b)
    e = keras.layers.Dropout(0.5, name='dropout')(c)
    model = keras.models.Model([a, b], [d, e])
    return model
  
  
  def strategy_minus_tpu_combinations():
    return tf.__internal__.test.combinations.combine(
        distribution=strategies_minus_tpu, mode=['graph', 'eager'])
  
  
  def tpu_strategy_combinations():
    return tf.__internal__.test.combinations.combine(
        distribution=tpu_strategies, mode=['graph', 'eager'])
  
  
  def tpu_strategy_combinations_graph_only():
    return tf.__internal__.test.combinations.combine(distribution=tpu_strategies, mode=['graph'])
  
  
  def multi_worker_strategy_combinations_eager_only():
    return tf.__internal__.test.combinations.combine(
        distribution=multi_worker_mirrored_strategies, mode=['eager'])
  
  
  def all_strategy_combinations():
    return strategy_minus_tpu_combinations() + tpu_strategy_combinations(
    ) + multi_worker_strategy_combinations_eager_only()
  
  
  def all_strategy_minus_default_and_tpu_combinations():
    return tf.__internal__.test.combinations.combine(
        distribution=[
            tf.__internal__.distribute.combinations.one_device_strategy,
            tf.__internal__.distribute.combinations.one_device_strategy_gpu,
            tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
            tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus,
        ],
        mode=['graph', 'eager'])
  
  
  def all_strategy_combinations_minus_default():
    return (all_strategy_minus_default_and_tpu_combinations() +
            tpu_strategy_combinations() +
            multi_worker_strategy_combinations_eager_only())
  
  
  def strategy_and_optimizer_combinations():
    non_tpu_strategies = tf.__internal__.test.combinations.times(
        strategy_minus_tpu_combinations(),
        tf.__internal__.test.combinations.combine(
            optimizer=[
                optimizer_combinations.adagrad_optimizer_v1_fn,
                optimizer_combinations.adam_optimizer_v1_fn,
                optimizer_combinations.gradient_descent_optimizer_v1_fn,
                optimizer_combinations.rmsprop_optimizer_v1_fn,
                optimizer_combinations.adadelta_optimizer_keras_v2_fn,
                optimizer_combinations.adagrad_optimizer_keras_v2_fn,
                optimizer_combinations.adam_optimizer_keras_v2_fn,
                optimizer_combinations.adamax_optimizer_keras_v2_fn,
                optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
                optimizer_combinations.nadam_optimizer_keras_v2_fn,
                optimizer_combinations.rmsprop_optimizer_keras_v2_fn,
                optimizer_combinations.ftrl_optimizer_keras_v2_fn
            ]))
    tpu_strategies_graph = tf.__internal__.test.combinations.combine(
        distribution=tpu_strategies,
        mode=['graph'],
        optimizer=[
            optimizer_combinations.adagrad_optimizer_v1_fn,
            optimizer_combinations.adam_optimizer_v1_fn,
            optimizer_combinations.gradient_descent_optimizer_v1_fn,
            optimizer_combinations.rmsprop_optimizer_v1_fn,
            optimizer_combinations.adagrad_optimizer_keras_v2_fn,
            optimizer_combinations.adam_optimizer_keras_v2_fn,
            optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
            optimizer_combinations.rmsprop_optimizer_keras_v2_fn
        ])
    tpu_strategies_eager = tf.__internal__.test.combinations.combine(
        distribution=tpu_strategies,
        mode=['eager'],
        optimizer=[
            optimizer_combinations.adagrad_optimizer_keras_v2_fn,
            optimizer_combinations.adam_optimizer_keras_v2_fn,
            optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
            optimizer_combinations.rmsprop_optimizer_keras_v2_fn
        ])
    multi_worker_eager = tf.__internal__.test.combinations.combine(
        distribution=multi_worker_mirrored_strategies,
        mode=['eager'],
        optimizer=[
            optimizer_combinations.adadelta_optimizer_keras_v2_fn,
            optimizer_combinations.adagrad_optimizer_keras_v2_fn,
            optimizer_combinations.adam_optimizer_keras_v2_fn,
            optimizer_combinations.adamax_optimizer_keras_v2_fn,
            optimizer_combinations.gradient_descent_optimizer_keras_v2_fn,
            optimizer_combinations.nadam_optimizer_keras_v2_fn,
            optimizer_combinations.rmsprop_optimizer_keras_v2_fn,
            optimizer_combinations.ftrl_optimizer_keras_v2_fn
        ])
    return (non_tpu_strategies + tpu_strategies_eager + tpu_strategies_graph +
            multi_worker_eager)
  
  
  class BatchCountingCB(keras.callbacks.Callback):
  
    def __init__(self):
      super(BatchCountingCB, self).__init__()
      self.train_begin_batches = []
      self.train_end_batches = []
      self.test_begin_batches = []
      self.test_end_batches = []
      self.predict_begin_batches = []
      self.predict_end_batches = []
  
    def on_train_batch_begin(self, batch, logs=None):
      self.train_begin_batches.append(batch)
  
    def on_train_batch_end(self, batch, logs=None):
      self.train_end_batches.append(batch)
  
    def on_test_batch_begin(self, batch, logs=None):
      self.test_begin_batches.append(batch)
  
    def on_test_batch_end(self, batch, logs=None):
      self.test_end_batches.append(batch)
  
    def on_predict_batch_begin(self, batch, logs=None):
      self.predict_begin_batches.append(batch)
  
    def on_predict_batch_end(self, batch, logs=None):
      self.predict_end_batches.append(batch)
  
  
  class TestDistributionStrategyWithNumpyArrays(tf.test.TestCase,
                                                parameterized.TestCase):
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_calculating_input_params_no_steps_no_batch_size(self, distribution):
      # Calculate the per_replica_batch_size scaling factor for strategies
      # that use per_core_batch_size
      replica_scale_factor = 1.0
      if not distributed_training_utils.global_batch_size_supported(distribution):
        replica_scale_factor = distribution.num_replicas_in_sync
  
      with self.cached_session():
        # Default global batch size 32 for input with 64 samples run in 2 steps
        steps, batch_size = distributed_training_utils_v1.get_input_params(
            distribution, 64, steps=None, batch_size=None)
        self.assertEqual(batch_size, 32 // replica_scale_factor)
        self.assertEqual(steps, 2)
  
        # Computed global batch size 20 is lower than 32 if we pass less samples.
        steps, batch_size = distributed_training_utils_v1.get_input_params(
            distribution, 20, steps=None, batch_size=None)
        self.assertEqual(batch_size, 20 // replica_scale_factor)
        self.assertEqual(steps, 1)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_calculating_input_params_with_steps_no_batch_size(
        self, distribution):
      # Calculate the per_replica_batch_size scaling factor for strategies
      # that use per_core_batch_size
      replica_scale_factor = 1.0
      if not distributed_training_utils.global_batch_size_supported(distribution):
        replica_scale_factor = distribution.num_replicas_in_sync
  
      with self.cached_session():
        # Computed global batch size is correct for number of specified 1 step
        steps, batch_size = distributed_training_utils_v1.get_input_params(
            distribution, 64, steps=1, batch_size=None)
        self.assertEqual(batch_size, 64 // replica_scale_factor)
        self.assertEqual(steps, 1)
  
        # Computed global batch size is correct for number of specified 2 steps
        steps, batch_size = distributed_training_utils_v1.get_input_params(
            distribution, 64, steps=2, batch_size=None)
        self.assertEqual(batch_size, 32 // replica_scale_factor)
        self.assertEqual(steps, 2)
  
        # All samples can not be consumed in specified number of steps
        with self.assertRaisesRegex(ValueError, 'not divisible by steps'):
          distributed_training_utils_v1.get_input_params(
              distribution, 63, steps=2, batch_size=None)
  
        # This cases is different for different strategies due to the
        # difference in supported batch size being global or per-replica.
        if replica_scale_factor == 1:
          # Computed global batch size is correct even if not sharadable
          steps, batch_size = distributed_training_utils_v1.get_input_params(
              distribution, 63, steps=3, batch_size=None)
          self.assertEqual(batch_size, 21)
          self.assertEqual(steps, 3)
        else:
          # Computed global batch size can not be sharded across replicas
          with self.assertRaisesRegex(
              ValueError, 'could not be sharded evenly '
              'across the sync replicas'):
            distributed_training_utils_v1.get_input_params(
                distribution, 63, steps=1, batch_size=None)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_calculating_input_params_no_steps_with_batch_size(
        self, distribution):
      # Calculate the per_replica_batch_size scaling factor for strategies
      # that use per_core_batch_size
      replica_scale_factor = 1.0
      if not distributed_training_utils.global_batch_size_supported(distribution):
        replica_scale_factor = distribution.num_replicas_in_sync
  
      with self.cached_session():
        # Computed steps is correct for specified batch size
        steps, batch_size = distributed_training_utils_v1.get_input_params(
            distribution, 64, steps=None, batch_size=16)
        self.assertEqual(batch_size, 16)
        self.assertEqual(steps, 4 // replica_scale_factor)
  
        # Computed steps is correct for specified batch size
        steps, batch_size = distributed_training_utils_v1.get_input_params(
            distribution, 64, steps=None, batch_size=32)
        self.assertEqual(batch_size, 32)
        self.assertEqual(steps, 2 // replica_scale_factor)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_calculating_input_params_with_steps_with_batch_size(
        self, distribution):
      with self.cached_session():
        # No change to steps and batch size if both specified and feasible
        steps, batch_size = distributed_training_utils_v1.get_input_params(
            distribution, 64, steps=5, batch_size=3)
        self.assertEqual(batch_size, 3)
        self.assertEqual(steps, 5)
  
        # Number of samples is less than global batch size * steps
        with self.assertRaisesRegex(ValueError, 'less than samples required'):
          distributed_training_utils_v1.get_input_params(
              distribution, 64, steps=10, batch_size=13)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_calling_model_with_numpy_arrays(self, distribution):
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(0.001)
          model = get_model()
          loss = 'mse'
          metrics = ['mae']
          model.compile(
              optimizer,
              loss,
              metrics=metrics)
  
          inputs = np.zeros((64, 3), dtype=np.float32)
          targets = np.zeros((64, 4), dtype=np.float32)
  
          # Call fit with validation data
          model.fit(
              inputs,
              targets,
              epochs=1,
              batch_size=2,
              verbose=0,
              validation_data=(inputs, targets))
  
          # TODO(anjalisridhar): We need tests for when the batch size and steps
          # are smaller and results in a 0 batch_size and steps value.
          model.evaluate(inputs, targets)
          model.evaluate(inputs, targets, batch_size=8)
  
          model.predict(inputs)
          model.predict(inputs, batch_size=8)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_calling_model_with_mixed_precision(self, distribution):
      if isinstance(distribution,
                    (tf.compat.v1.distribute.experimental.ParameterServerStrategy,
                     tf.distribute.experimental.ParameterServerStrategy,
                     tf.distribute.experimental.CentralStorageStrategy,
                     tf.compat.v1.distribute.experimental.CentralStorageStrategy)):
        self.skipTest('b/152097775')
      if backend.is_tpu_strategy(distribution):
        policy_name = 'mixed_bfloat16'
      else:
        policy_name = 'mixed_float16'
      with self.cached_session(), \
           distribution.scope(), \
           policy.policy_scope(policy_name):
        optimizer_fn = gradient_descent_keras.SGD
        optimizer = optimizer_fn(0.001)
        x = keras.layers.Input(shape=(3,), name='input')
        y = keras.layers.Dense(4, name='dense')(x)
        y = keras.layers.Activation('softmax', dtype='float32')(y)
        model = keras.Model(x, y)
        loss = 'mse'
        metrics = ['mae']
        model.compile(
            optimizer,
            loss,
            metrics=metrics)
  
        # We need to pass float32 since TPUs do not support float64, even though
        # these arrays will immediately be casted to bfloat16 on TPUs. We also
        # cannot pass bfloat16, as Numpy does not support it.
        inputs = np.zeros((64, 3), dtype='float32')
        targets = np.zeros((64, 4), dtype='float32')
  
        model.fit(
            inputs,
            targets,
            epochs=1,
            batch_size=2,
            verbose=0,
            validation_data=(inputs, targets))
  
        model.evaluate(inputs, targets)
        model.evaluate(inputs, targets, batch_size=8)
  
        model.predict(inputs)
        model.predict(inputs, batch_size=8)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_operator_overload_mixed_precision(self, distribution):
      # Regression test that tests a fixed bug does not reoccur. Adding an
      # AutoCastVariable to a tensor on a TPU, where the variable was the LHS of
      # the '+' operator, used to cause the gradient w.r.t. the variable to be
      # None.
      if isinstance(distribution,
                    (tf.compat.v1.distribute.experimental.ParameterServerStrategy,
                     tf.distribute.experimental.ParameterServerStrategy,
                     tf.distribute.experimental.CentralStorageStrategy,
                     tf.compat.v1.distribute.experimental.CentralStorageStrategy)):
        self.skipTest('b/152097775')
  
      if backend.is_tpu_strategy(distribution):
        policy_name = 'mixed_bfloat16'
      else:
        policy_name = 'mixed_float16'
  
      class MyLayer(keras.layers.Layer):
  
        def build(self, _):
          self.v1 = self.add_weight('v', ())
          self.v2 = self.add_weight('v', ())
  
        def call(self, inp):
          inp += self.v1
          return self.v2 + inp
  
      with self.cached_session(), distribution.scope():
        layer = MyLayer(dtype=policy_name)
        def run_fn():
          x = np.array([1.])
          with tf.GradientTape() as tape:
            y = layer(x)
          grad_v1, grad_v2 = tape.gradient(y, [layer.v1, layer.v2])
          return grad_v1, grad_v2
        if tf.executing_eagerly():
          run_fn = tf.function(run_fn)
  
        grad_v1, grad_v2 = distribution.run(run_fn)
        self.assertIsNotNone(grad_v1)
        self.assertIsNotNone(grad_v2)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=[tf.__internal__.distribute.combinations.one_device_strategy],
            mode=['graph', 'eager']))
    def test_optimizer_in_cross_replica_context_raises_error(self, distribution):
  
      with self.cached_session(), distribution.scope():
        model = keras.models.Sequential([keras.layers.Dense(1)])
        x = np.array([[1.]])
        with tf.GradientTape() as tape:
          y = model(x)
        gradients = tape.gradient(y, model.trainable_variables)
        optimizer = gradient_descent_keras.SGD()
  
        with self.assertRaisesRegex(RuntimeError,
                                    'cannot be called in cross-replica context'):
          optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_calling_model_with_nested_numpy_arrays(self, distribution):
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(learning_rate=0.001)
          model = multi_input_output_model()
          loss = 'mse'
          model.compile(
              optimizer,
              loss)
  
        input_a_np = np.asarray(np.random.random((64, 3)), dtype=np.float32)
        input_b_np = np.asarray(np.random.random((64, 5)), dtype=np.float32)
        inputs = [input_a_np, input_b_np]
  
        output_d_np = np.asarray(np.random.random((64, 7)), dtype=np.float32)
        output_e_np = np.asarray(np.random.random((64, 7)), dtype=np.float32)
        targets = [output_d_np, output_e_np]
  
        # Call fit with validation data
        model.fit(inputs, targets, epochs=1, batch_size=8, verbose=0)
  
        # TODO(anjalisridhar): We need tests for when the batch size and steps are
        # smaller and results in a 0 batch_size and steps value.
        model.evaluate(inputs, targets)
        model.evaluate(inputs, targets, batch_size=8)
  
        model.predict(inputs)
        model.predict(inputs, batch_size=8)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=strategies_minus_tpu, mode=['graph', 'eager']) +
        tf.__internal__.test.combinations.combine(
            distribution=multi_worker_mirrored_strategies, mode=['eager']))
    def test_numpy_with_sample_weights(self, distribution):
      with self.cached_session(), distribution.scope():
        model = get_sample_weights_model()
        optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate=0.001)
        loss = 'mse'
        model.compile(
            optimizer,
            loss)
  
        inputs = np.array([[0], [1], [2], [3]], np.float32)
        targets = np.array([[2], [4], [6], [8]], np.float32)
        sample_weights = np.array([0.25, 0.5, 0.75, 1], np.float32)
  
        result = model.evaluate(
            inputs,
            targets,
            batch_size=2,
            sample_weight=sample_weights,
            verbose=1)
  
        # The per sample loss is multiplied by the corresponding sample weight.
        # The average of these weighted losses is the return value of the
        # `evaluate` call. For example, in the test above the average weighted
        # loss is calculated in the following manner:
  
        # batch_1 = (((2-0)^2) * 0.25 + ((4-1)^2) * 0.5) / 2 = 5.5 / 2 = 2.75
        # batch_2 = (((6-2)^2 * 0.75) + ((8-3)^2 * 1)) / 2 = 37 / 2 = 18.5
        # final result = (batch_1 + batch_2) / 2 = 10.625.
        # The first time we divide by number of input samples and the second time
        # we divide by number of steps/batches that the loss is aggregated over.
        self.assertAllClose(result, 10.625)
  
        # We now test without passing sample_weights:
        # batch_1 = ((2-0)^2) + ((4-1)^2) / 2 = 13 / 2 = 6.5
        # batch_2 = ((6-2)^2) + ((8-3)^2) / 2 = 41 / 2 = 20.5
        # final result = (batch_1 + batch_2) / 2 =  27 / 2 = 13.5
        result = model.evaluate(inputs, targets, batch_size=2, verbose=1)
        self.assertAllClose(result, 13.5)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_flatten_predict_outputs(self, distribution):
      with self.cached_session():
        with distribution.scope():
          model = multi_input_output_model()
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(learning_rate=0.001)
          loss = 'mse'
          model.compile(
              optimizer,
              loss)
  
        # We take 6 input samples with each input having a dimension of 3 or 5.
        input_a_np = np.asarray(np.random.random((6, 3)), dtype=np.float32)
        input_b_np = np.asarray(np.random.random((6, 5)), dtype=np.float32)
        inputs = [input_a_np, input_b_np]
  
        outs = model.predict(inputs)
        # `predict` a list that is equal in length to the number of model outputs.
        # In this test our model has two outputs and each element of `outs`
        # corresponds to all the samples of one of the model outputs.
        self.assertLen(outs, 2)
        # Each of the output samples have a dimension of 7. We should process all
        # the available input samples(6).
        self.assertAllEqual([6, 7], outs[0].shape)
        self.assertAllEqual([6, 7], outs[1].shape)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(tpu_strategy_combinations_graph_only(),
                           tf.__internal__.test.combinations.combine(batch_size=[4, 6])))
    def test_evaluate_with_partial_batch(self, distribution, batch_size):
      with self.cached_session():
        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
        loss = 'mse'
        metrics = ['mae', keras.metrics.CategoricalAccuracy()]
  
        with distribution.scope():
          model_with_ds_strategy = get_model()
          model_with_ds_strategy.compile(optimizer, loss, metrics=metrics)
  
        cpu_model = get_model()
        cpu_model.compile(optimizer, loss, metrics=metrics)
  
        x = np.random.random((10, 3)).astype('float32')
        y = np.random.random((10, 4)).astype('float32')
  
        # As sample size is 10, we batch by 4 so that the last batch is
        # a partial batch. Also `evaluate()` using numpy array as inputs without
        # distribution strategy uses entire sample as a single batch. As so,
        # we remove parameters `batch_size` and `steps`.
        cpu_model.set_weights(model_with_ds_strategy.get_weights())
        evaluate_ground_truth = cpu_model.evaluate(x, y)
  
        # We don't compare the loss as loss is currently not computed as metric
        # in Keras, the loss value is inaccurate for last partial batch due to
        # more weights for the last batch samples.
        steps = np.ceil(10.0 / batch_size)
        self.assertAllClose(
            model_with_ds_strategy.evaluate(
                x, y, batch_size=batch_size, steps=steps)[1:],
            evaluate_ground_truth[1:],
            atol=1e-5,
            rtol=1e-5)
        # Test that `steps` is inferred correctly when final partial batch exists.
        self.assertAllClose(
            model_with_ds_strategy.evaluate(x, y, batch_size=batch_size)[1:],
            evaluate_ground_truth[1:],
            atol=1e-5,
            rtol=1e-5)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(
            tpu_strategy_combinations_graph_only()))
    def test_predict_with_partial_batch(self, distribution):
      with self.cached_session():
        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
        loss = 'mse'
  
        with distribution.scope():
          model_with_ds_strategy = get_model()
          model_with_ds_strategy.compile(
              optimizer,
              loss)
  
        cpu_model = get_model()
        cpu_model.compile(optimizer, loss)
  
        inputs = np.random.random((10, 3)).astype(np.float32)
  
        # As sample size is 10, we batch by 4 so that the last batch is
        # a partial batch. Also `predict()` using numpy array as inputs without
        # distribution strategy uses entire sample as a single batch. As so,
        # we remove parameters `batch_size` and `steps`.
        cpu_model.set_weights(model_with_ds_strategy.get_weights())
        predict_ground_truth = cpu_model.predict(inputs)
        self.assertAllClose(
            model_with_ds_strategy.predict(inputs, batch_size=4, steps=3),
            predict_ground_truth,
            atol=1e-5,
            rtol=1e-5)
        # Test that `steps` is inferred correctly when final partial batch exists.
        self.assertAllClose(
            model_with_ds_strategy.predict(inputs, batch_size=4),
            predict_ground_truth,
            atol=1e-5,
            rtol=1e-5)
  
    @tf.__internal__.distribute.combinations.generate(tpu_strategy_combinations_graph_only())
    def test_no_target_model(self, distribution):
      with self.cached_session():
        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
  
        class MyLayer(keras.layers.Layer):
  
          def call(self, inputs, training=None):
            self.add_loss(tf.reduce_sum(inputs), inputs=True)
            return inputs
  
        with distribution.scope():
          model = keras.models.Sequential()
          model.add(
              keras.layers.Dense(16, activation='relu', input_shape=_INPUT_SIZE))
          model.add(MyLayer())
          model.add(keras.layers.Dense(_NUM_CLASS, activation='softmax'))
  
          model.compile(optimizer)
          inputs = np.zeros((20, 10), np.float32)
  
          model.fit(inputs, epochs=1, steps_per_epoch=2)
          model.predict(inputs, steps=1)
          model.evaluate(inputs, steps=1)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(
            tpu_strategy_combinations_graph_only()))
    def test_predict_multi_output_model_with_partial_batch(
        self, distribution):
      with self.cached_session():
        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
        loss = 'mse'
  
        with distribution.scope():
          model_with_ds_strategy = simple_multi_inputs_multi_outputs_model()
          model_with_ds_strategy.compile(
              optimizer,
              loss)
  
        cpu_model = simple_multi_inputs_multi_outputs_model()
        cpu_model.compile(optimizer, loss)
  
        input_data, _ = get_multi_inputs_multi_outputs_data()
        input_dict = {
            'input_a': input_data['input_a'],
            'input_b': input_data['input_b'],
        }
  
        # As sample size is 200, we batch by 18 so that the last batch is
        # a partial batch. Also `fit()` using numpy array as inputs without
        # distribution strategy uses entire sample as a single batch. As so,
        # we remove parameters `batch_size` and `steps`.
        cpu_model.set_weights(model_with_ds_strategy.get_weights())
        self.assertAllClose(
            model_with_ds_strategy.predict(input_dict, batch_size=18, steps=12),
            cpu_model.predict(input_dict),
            atol=1e-4,
            rtol=1e-4)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_gradients_are_none(self, distribution):
  
      if not tf.executing_eagerly():
        self.skipTest('None gradients are not supported in graph mode')
  
      class DenseWithExtraWeight(keras.layers.Dense):
  
        def build(self, input_shape):
          # Gradients w.r.t. extra_weights are None
          self.extra_weight_1 = self.add_weight('extra_weight_1', shape=(),
                                                initializer='ones')
          super(DenseWithExtraWeight, self).build(input_shape)
          self.extra_weight_2 = self.add_weight('extra_weight_2', shape=(),
                                                initializer='ones')
  
      with distribution.scope():
        model = keras.Sequential([DenseWithExtraWeight(4, input_shape=(4,))])
        model.compile('adam', 'mse')
  
      inputs = np.random.normal(size=(64, 4))
      targets = np.random.normal(size=(64, 4))
      old_kernel = model.get_weights()[1]
      model.fit(inputs, targets)
      new_kernel = model.get_weights()[1]
      self.assertNotAllEqual(old_kernel, new_kernel)
  
  
  class TestDistributionStrategyWithDatasets(tf.test.TestCase,
                                             parameterized.TestCase):
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_calling_model_on_same_dataset(self, distribution):
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(0.001)
          model = get_model()
          loss = 'mse'
          metrics = ['mae', keras.metrics.CategoricalAccuracy()]
          model.compile(
              optimizer,
              loss,
              metrics=metrics)
  
        dataset = get_dataset(distribution)
  
        # Call fit with validation data
        model.fit(
            dataset,
            epochs=1,
            steps_per_epoch=2,
            verbose=0,
            validation_data=dataset,
            validation_steps=2)
        model.fit(
            dataset,
            epochs=1,
            steps_per_epoch=2,
            verbose=0,
            validation_data=dataset,
            validation_steps=2)
        model.predict(get_predict_dataset(distribution), steps=2)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_model_interleaved_eval_same_as_direct_eval(
        self, distribution):
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          user_controlled_model = get_model()
          user_controlled_model.compile(
              optimizer_fn(0.001),
              loss='mse',
              metrics=['mae', keras.metrics.CategoricalAccuracy()])
  
          interleaved_model = get_model()
          interleaved_model.set_weights(user_controlled_model.get_weights())
          interleaved_model.compile(
              optimizer_fn(0.001),
              loss='mse',
              metrics=['mae', keras.metrics.CategoricalAccuracy()])
  
        dataset = get_dataset(distribution)
  
        # Call fit with validation interleaved
        interleaved_output = interleaved_model.fit(
            dataset,
            epochs=2,
            steps_per_epoch=2,
            verbose=1,
            validation_data=dataset,
            validation_steps=2,
            shuffle=False)
  
        # Manually control the validation running after each epoch.
        user_controlled_output = []
        for _ in range(2):
          user_controlled_model.fit(
              dataset, epochs=1, steps_per_epoch=2, verbose=1, shuffle=False)
          user_controlled_output.append(
              user_controlled_model.evaluate(dataset, steps=2))
  
        self.assertEqual(interleaved_output.history['val_loss'],
                         [x[0] for x in user_controlled_output])
        val_mean_absolute_error = interleaved_output.history.get(
            'val_mean_absolute_error')
        if not val_mean_absolute_error:
          # The name of the metric changed in TF2.0
          val_mean_absolute_error = interleaved_output.history['val_mae']
        self.assertEqual(val_mean_absolute_error,
                         [x[1] for x in user_controlled_output])
        self.assertEqual(interleaved_output.history['val_categorical_accuracy'],
                         [x[2] for x in user_controlled_output])
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_fit_with_tuple_and_dict_dataset_inputs(self, distribution):
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(learning_rate=0.001)
          model = multi_input_output_model()
          loss = 'mse'
          metrics = ['mae', keras.metrics.CategoricalAccuracy()]
          model.compile(
              optimizer,
              loss,
              metrics=metrics)
  
        input_a_np = np.random.random((10, 3)).astype('float32')
        input_b_np = np.random.random((10, 5)).astype('float32')
        output_d_np = np.random.random((10, 7)).astype('float32')
        output_e_np = np.random.random((10, 7)).astype('float32')
  
        # Test with tuples
        dataset_tuple = tf.data.Dataset.from_tensor_slices(
            ((input_a_np, input_b_np), (output_d_np, output_e_np)))
        dataset_tuple = dataset_tuple.repeat(100)
        dataset_tuple = dataset_tuple.batch(10)
  
        model.fit(dataset_tuple, epochs=1, steps_per_epoch=2, verbose=1)
  
        # Test with dict
        dataset_dict = tf.data.Dataset.from_tensor_slices(({
            'input_a': input_a_np,
            'input_b': input_b_np
        }, (output_d_np, output_e_np)))
        dataset_dict = dataset_dict.repeat(100)
        dataset_dict = dataset_dict.batch(10)
  
        model.fit(dataset_dict, epochs=1, steps_per_epoch=2, verbose=1)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_fit_with_dictionary_in_the_dataset_b135161171(
        self, distribution):
  
      if backend.is_tpu_strategy(distribution):
        self.skipTest('b/142805125')
  
      def custom_loss(predict, label, weight):
        bce = keras.losses.binary_crossentropy(label, predict)
        return tf.reduce_mean(bce * weight)
  
      with self.cached_session():
        with distribution.scope():
          input_img = keras.layers.Input([64, 64, 3], name='img')
          input_lbl = keras.layers.Input([64, 64, 1], name='lbl')
          input_weight = keras.layers.Input([64, 64], name='weight')
          predict = keras.layers.Conv2D(2, [1, 1], padding='same')(input_img)
          loss_lambda = keras.layers.Lambda(
              lambda x: custom_loss(*x), name='my_loss')
          my_loss = loss_lambda([predict, input_lbl, input_weight])
          model = keras.models.Model(
              inputs=[input_img, input_lbl, input_weight],
              outputs=[predict, my_loss])
          model.add_loss(model.get_layer('my_loss').output)
          model.compile(
              optimizer='adam')
  
        if tf.executing_eagerly():
  
          def map_fn(img, lbl, weight):
            inputs = {'img': img, 'lbl': lbl, 'weight': weight}
            return (inputs,)
        else:
  
          def map_fn(img, lbl, weight):
            inputs = {'img': img, 'lbl': lbl, 'weight': weight}
            return inputs, {}
  
        fake_imgs = np.ones([50, 64, 64, 3], dtype=np.float32)
        fake_lbls = np.ones([50, 64, 64, 1], dtype=np.float32)
        fake_weights = np.ones([50, 64, 64], dtype=np.float32)
  
        data = tf.data.Dataset.from_tensor_slices(
            (fake_imgs, fake_lbls, fake_weights)).map(map_fn).batch(10)
  
        model.fit(data)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_fit_eval_and_predict_methods_on_dataset_without_steps(
        self, distribution):
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(0.001)
          model = get_model()
          loss = 'mse'
          metrics = ['mae', keras.metrics.CategoricalAccuracy()]
          model.compile(
              optimizer,
              loss,
              metrics=metrics)
  
        inputs = np.zeros((1000, 3), dtype=np.float32)
        targets = np.zeros((1000, 4), dtype=np.float32)
        # steps/steps_per_epoch are calculated when using numpy arrays as
        # input data.
        fit_with_numpy = model.fit(
            inputs, targets, epochs=1, batch_size=10).history
        eval_with_numpy = model.evaluate(inputs, targets, batch_size=10)
        predict_with_numpy = model.predict(inputs, batch_size=10)
  
        dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
        dataset = dataset.batch(10, drop_remainder=True)
        fit_with_ds = model.fit(dataset, epochs=1).history
        eval_with_ds = model.evaluate(dataset)
        predict_dataset = tf.data.Dataset.from_tensor_slices(inputs)
        predict_dataset = predict_dataset.batch(10, drop_remainder=True)
        predict_with_ds = model.predict(predict_dataset)
        self.assertAllClose(fit_with_numpy, fit_with_ds, atol=1e-4, rtol=1e-4)
        self.assertAllClose(eval_with_numpy, eval_with_ds, atol=1e-4, rtol=1e-4)
        self.assertAllClose(
            predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_predict_on_dataset_with_unknown_cardinality_without_steps(
        self, distribution, mode):
  
      if mode == 'graph' and backend.is_tpu_strategy(distribution):
        self.skipTest('partial batch not supported with TPU in graph mode.')
  
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(0.001)
          model = get_model()
          loss = 'mse'
          metrics = ['mae', keras.metrics.CategoricalAccuracy()]
          model.compile(optimizer, loss, metrics=metrics)
  
        inputs = np.zeros((20, 3), dtype=np.float32)
        # steps/steps_per_epoch are calculated when using numpy arrays as
        # input data.
        predict_with_numpy = model.predict(inputs, batch_size=10)
  
        predict_dataset = convert_numpy_to_dataset_with_unknown_cardinality(
            inputs)
  
        self.assertEqual(
            keras.backend.get_value(tf.data.experimental.cardinality(predict_dataset)),
            tf.data.experimental.UNKNOWN_CARDINALITY)
  
        predict_with_ds = model.predict(predict_dataset)
        self.assertAllClose(
            predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_on_dataset_with_unknown_cardinality_without_steps(
        self, distribution, mode):
      # TODO(b/155867206): Investigate why this test occasionally segfaults on TPU
      # in eager mode.
      if mode == 'eager' and backend.is_tpu_strategy(distribution):
        self.skipTest('caused segfault with TPU in eager mode.')
  
      if mode == 'graph' and backend.is_tpu_strategy(distribution):
        self.skipTest('partial batch not supported with TPU in graph mode.')
  
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(0.001)
          model = get_model()
          loss = 'mse'
          metrics = ['mae', keras.metrics.CategoricalAccuracy()]
          model.compile(
              optimizer,
              loss,
              metrics=metrics)
  
        inputs = np.zeros((100, 3), dtype=np.float32)
        targets = np.zeros((100, 4), dtype=np.float32)
        # steps/steps_per_epoch are calculated when using numpy arrays as
        # input data.
        fit_with_numpy = model.fit(
            inputs, targets, epochs=1, batch_size=10).history
        fit_with_numpy_multiple_epochs = model.fit(
            inputs, targets, epochs=2, batch_size=10).history
        eval_with_numpy = model.evaluate(inputs, targets, batch_size=10)
        predict_with_numpy = model.predict(inputs, batch_size=10)
  
        dataset = convert_numpy_to_dataset_with_unknown_cardinality(
            inputs, targets)
        predict_dataset = convert_numpy_to_dataset_with_unknown_cardinality(
            inputs)
  
        self.assertEqual(
            keras.backend.get_value(tf.data.experimental.cardinality(dataset)),
            tf.data.experimental.UNKNOWN_CARDINALITY)
        self.assertEqual(
            keras.backend.get_value(tf.data.experimental.cardinality(predict_dataset)),
            tf.data.experimental.UNKNOWN_CARDINALITY)
  
        eval_with_ds = model.evaluate(dataset)
        predict_with_ds = model.predict(predict_dataset)
        self.assertAllClose(eval_with_numpy, eval_with_ds, atol=1e-4, rtol=1e-4)
        self.assertAllClose(
            predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4)
  
        fit_with_ds = model.fit(dataset, epochs=1).history
        fit_with_ds_multiple_epochs = model.fit(dataset, epochs=2).history
        self.assertAllClose(fit_with_numpy, fit_with_ds, atol=1e-4, rtol=1e-4)
        self.assertAllClose(
            fit_with_numpy_multiple_epochs,
            fit_with_ds_multiple_epochs,
            atol=1e-4,
            rtol=1e-4)
  
    @tf.__internal__.distribute.combinations.generate(tpu_strategy_combinations_graph_only())
    def test_on_dataset_with_unknown_cardinality(self, distribution):
      with self.cached_session():
        with distribution.scope():
          model = get_model()
          loss = 'mse'
          metrics = ['mae', keras.metrics.CategoricalAccuracy()]
          model.compile(
              tf.compat.v1.train.GradientDescentOptimizer(0.001),
              loss,
              metrics=metrics)
  
        inputs = np.zeros((1000, 3), dtype=np.float32)
        targets = np.zeros((1000, 4), dtype=np.float32)
        # steps/steps_per_epoch are calculated when using numpy arrays as
        # input data.
        eval_with_numpy = model.evaluate(inputs, targets, batch_size=10)
        predict_with_numpy = model.predict(inputs, batch_size=10)
  
        dataset = convert_numpy_to_dataset_with_unknown_cardinality(
            inputs, targets)
        predict_dataset = convert_numpy_to_dataset_with_unknown_cardinality(
            inputs)
  
        self.assertEqual(
            keras.backend.get_value(tf.data.experimental.cardinality(dataset)),
            tf.data.experimental.UNKNOWN_CARDINALITY)
        self.assertEqual(
            keras.backend.get_value(tf.data.experimental.cardinality(predict_dataset)),
            tf.data.experimental.UNKNOWN_CARDINALITY)
  
        eval_with_ds = model.evaluate(dataset, steps=100)
        predict_with_ds = model.predict(predict_dataset, steps=100)
        self.assertAllClose(eval_with_numpy, eval_with_ds, atol=1e-4, rtol=1e-4)
        self.assertAllClose(
            predict_with_numpy, predict_with_ds, atol=1e-4, rtol=1e-4)
  
        with self.assertRaisesRegex(ValueError,
                                    'Number of steps could not be inferred'):
          model.fit(dataset, epochs=1)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_fit_eval_and_predict_methods_on_dataset(
        self, distribution):
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(0.001)
          model = get_model()
          loss = 'mse'
          metrics = ['mae', keras.metrics.CategoricalAccuracy()]
          model.compile(
              optimizer,
              loss,
              metrics=metrics)
  
        dataset = get_dataset(distribution)
  
        model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
        model.evaluate(dataset, steps=2, verbose=1)
        model.predict(get_predict_dataset(distribution), steps=2)
  
    @tf.__internal__.distribute.combinations.generate(strategy_and_optimizer_combinations())
    def test_fit_eval_and_predict_with_optimizer(self, distribution, optimizer):
      with self.cached_session():
  
        with distribution.scope():
  
          model = get_model()
          loss = 'mse'
          model.compile(
              optimizer(),
              loss)
  
        dataset = get_dataset(distribution)
  
        model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
        model.evaluate(dataset, steps=2, verbose=1)
        model.predict(get_predict_dataset(distribution), steps=2)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=[
                tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
                tf.__internal__.distribute.combinations.one_device_strategy
            ],
            mode=['graph', 'eager']))
    def test_dataset_wrong_input_shape(self, distribution, mode):
      if mode == 'graph':
        self.skipTest(
            'TODO(b/120943676, b/120957836): Re-enable for graph once the '
            'validation code is restored.')
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(learning_rate=0.001)
          model = get_model()
          loss = 'mse'
          model.compile(
              optimizer,
              loss)
  
        # Wrong input shape
        inputs = np.zeros((10, 5), dtype=np.float32)
        targets = np.zeros((10, 4), dtype=np.float32)
        dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
        dataset = dataset.repeat(100)
        dataset = dataset.batch(10)
  
        with self.assertRaisesRegex(ValueError, 'is incompatible with'):
          model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=0)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=[
                tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu
            ],
            mode=['graph', 'eager']))
    def test_dataset_external_batch_input_validation(
        self, distribution):
      with self.cached_session():
        with distribution.scope():
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(learning_rate=0.001)
          model = get_model()
          loss = 'mse'
          model.compile(
              optimizer,
              loss)
  
        # Batching is done outside tf.data's `batch`
        inputs = np.zeros((100, 10, 3), dtype=np.float32)
        targets = np.zeros((100, 10, 4), dtype=np.float32)
        dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
        dataset = dataset.repeat(100)
  
        model.fit(dataset, epochs=1, steps_per_epoch=2, verbose=1)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=[
                tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
                tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus
            ],
            mode=['graph', 'eager']))
    def test_learning_phase_value(self, distribution):
      # TODO(anjalisridhar): Modify this test to use Lambdas since we can compare
      # meaningful values. Currently we don't pass the learning phase if the
      # Lambda layer uses the learning phase.
      with self.cached_session():
        with distribution.scope():
          x = keras.layers.Input(shape=(1,), name='input')
          y = keras.layers.Dense(1, kernel_initializer='ones')(x)
          z = keras.layers.Dropout(0.9999)(y)
          model = keras.Model(x, z)
          initial_weights = model.get_weights()
  
          optimizer_fn = gradient_descent_keras.SGD
          optimizer = optimizer_fn(0.005)
          loss = 'mse'
          metrics = ['acc']
          model.compile(
              optimizer,
              loss,
              metrics=metrics)
  
        batch_size = 8
        if isinstance(distribution, (tf.distribute.MirroredStrategy,
                                     tf.compat.v1.distribute.MirroredStrategy)):
          # MirroredStrategy uses global batch size.
          batch_size = 8 * distribution.num_replicas_in_sync
  
        inputs = np.ones((10, 1), dtype=np.float32)
        targets = np.ones((10, 1), dtype=np.float32)
        dataset = tf.data.Dataset.from_tensor_slices((inputs, targets))
        dataset = dataset.repeat().batch(batch_size)
        hist = model.fit(dataset, epochs=1, steps_per_epoch=20, verbose=1)
        self.assertAlmostEqual(hist.history['acc'][0], 0, 0)
  
        with distribution.scope():
          model.set_weights(initial_weights)
        # TODO(psv/anjalisridhar): Enable these lines after we fix b/117431185.
        # evaluate_output = model.evaluate(dataset, steps=20)
        # self.assertAlmostEqual(evaluate_output[1], 1, 0)
  
        inputs = np.ones((10, 1), dtype=np.float32)
        predict_dataset = tf.data.Dataset.from_tensor_slices(inputs)
  
        predict_dataset = predict_dataset.repeat().batch(batch_size)
        output = model.predict(predict_dataset, steps=10)
        # `predict` runs for 10 steps
        ref_output = np.ones((160, 1), dtype=np.float32)
        self.assertArrayNear(output, ref_output, 1e-1)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def testOptimizerWithCallbacks(self, distribution):
      with self.cached_session():
        with distribution.scope():
          model = get_model()
          optimizer = gradient_descent_keras.SGD(0.01)
          loss = 'mse'
          model.compile(
              optimizer,
              loss)
  
        dataset = get_dataset(distribution)
  
        def schedule(_):
          return 0.001
  
        model.fit(
            dataset,
            epochs=1,
            steps_per_epoch=2,
            verbose=0,
            callbacks=[keras.callbacks.LearningRateScheduler(schedule)])
        self.assertAllClose(0.001, keras.backend.get_value(model.optimizer.lr))
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(tpu_strategy_combinations_graph_only(),
                           tf.__internal__.test.combinations.combine(batch_size=[4, 6])))
    def test_evaluate_with_dataset_with_partial_batch(self, distribution,
                                                      batch_size):
      with self.cached_session():
        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
        loss = 'mse'
        metrics = ['mae', keras.metrics.CategoricalAccuracy()]
  
        with distribution.scope():
          model_with_ds_strategy = get_model()
          model_with_ds_strategy.compile(optimizer, loss, metrics=metrics)
  
        cpu_model = get_model()
        cpu_model.compile(optimizer, loss, metrics=metrics)
  
        x = np.random.random((10, 3)).astype('float32')
        y = np.random.random((10, 4)).astype('float32')
        dataset = tf.data.Dataset.from_tensor_slices((x, y))
  
        # As sample size is 10, we make the last batch a partial batch.
        cpu_model.set_weights(model_with_ds_strategy.get_weights())
        dataset_with_partial_batch = dataset.batch(batch_size)
  
        # We don't compare the loss as loss is currently not computed as metric
        # in Keras, the loss value is inaccurate for last partial batch due to
        # more weights for the last batch samples.
        steps = np.ceil(10.0 / batch_size)
        self.assertAllClose(
            model_with_ds_strategy.evaluate(
                dataset_with_partial_batch, steps=steps)[1:],
            cpu_model.evaluate(dataset_with_partial_batch, steps=steps)[1:],
            atol=1e-5,
            rtol=1e-5)
        self.assertAllClose(
            model_with_ds_strategy.evaluate(dataset_with_partial_batch)[1:],
            cpu_model.evaluate(dataset_with_partial_batch)[1:],
            atol=1e-5,
            rtol=1e-5)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(
            tpu_strategy_combinations_graph_only()))
    def test_predict_with_dataset_with_partial_batch(
        self, distribution):
      with self.cached_session():
        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
        loss = 'mse'
  
        with distribution.scope():
          model_with_ds_strategy = get_model()
          model_with_ds_strategy.compile(
              optimizer,
              loss)
  
        cpu_model = get_model()
        cpu_model.compile(optimizer, loss)
  
        inputs = np.random.random((10, 3)).astype(np.float32)
        dataset = tf.data.Dataset.from_tensor_slices((inputs))
  
        # As sample size is 10, we batch by 4 so that the last batch is
        # a partial batch.
        dataset_with_partial_batch = dataset.batch(4)
        cpu_model.set_weights(model_with_ds_strategy.get_weights())
  
        self.assertAllClose(
            model_with_ds_strategy.predict(dataset_with_partial_batch, steps=3),
            cpu_model.predict(dataset_with_partial_batch, steps=3),
            atol=1e-5,
            rtol=1e-5)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(
            tpu_strategy_combinations_graph_only()))
    def test_predict_multi_output_model_with_dataset_with_partial_batch(
        self, distribution):
      with self.cached_session():
        optimizer = tf.compat.v1.train.GradientDescentOptimizer(0.001)
        loss = 'mse'
  
        with distribution.scope():
          model_with_ds_strategy = simple_multi_inputs_multi_outputs_model()
          model_with_ds_strategy.compile(
              optimizer,
              loss)
  
        cpu_model = simple_multi_inputs_multi_outputs_model()
        cpu_model.compile(optimizer, loss)
  
        input_data, _ = get_multi_inputs_multi_outputs_data()
        input_dict = {
            'input_a': input_data['input_a'],
            'input_b': input_data['input_b'],
        }
  
        dataset = tf.data.Dataset.from_tensor_slices(input_dict)
  
        # As sample size is 200, we batch by 18 using 12 steps per epoch so
        # that the last batch is a partial batch.
        dataset_with_partial_batch = dataset.batch(18)
        cpu_model.set_weights(model_with_ds_strategy.get_weights())
  
        self.assertAllClose(
            model_with_ds_strategy.predict(dataset_with_partial_batch, steps=12),
            cpu_model.predict(dataset_with_partial_batch, steps=12),
            atol=1e-4,
            rtol=1e-4)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations_minus_default())
    def test_match_model_input_matches_with_dataset_tensors(self, distribution):
  
      def _create_model_input_output_tensors():
        input_a = keras.layers.Input(shape=(16,), name='z_input_sorted_last')
        input_b = keras.layers.Input(shape=(32,), name='a_input_sorted_first')
        intermediate_a = keras.layers.Dense(10)(input_a)
        intermediate_b = keras.layers.Dense(10)(input_b)
        merged = keras.layers.Add()([intermediate_a, intermediate_b])
        output = keras.layers.Dense(2)(merged)
        return input_a, input_b, output
  
      input_dict = {
          'z_input_sorted_last': np.random.rand(32, 16).astype(np.float32),
          'a_input_sorted_first': np.random.rand(32, 32).astype(np.float32)
      }
      target = np.ones((32, 2), dtype=np.float32)
      dataset = tf.data.Dataset.from_tensor_slices((input_dict, target))
      dataset = dataset.batch(4, drop_remainder=True)
  
      with self.cached_session():
        with distribution.scope():
          input_a, input_b, output = _create_model_input_output_tensors()
          # `input_a`, which has input name that comes last in alphanumeric
          # order, is the first input of the model input layers. If tensors
          # from `input_dict` is blindly flattened and passed to model
          # inputs incorrectly, this would result in `input_a` input layer
          # matching with tensor `a_input_sorted_first` and would result in
          # shape mismatch.
          model_with_array_input = keras.models.Model(
              inputs=[input_a, input_b], outputs=output)
          model_with_array_input.compile('sgd', 'mse')
          model_weights = model_with_array_input.get_weights()
          model_with_array_input_fit = model_with_array_input.fit(
              dataset, steps_per_epoch=1, epochs=1).history
  
          input_a, input_b, output = _create_model_input_output_tensors()
          model_with_dict_input = keras.models.Model(
              inputs={
                  'z_input_sorted_last': input_a,
                  'a_input_sorted_first': input_b,
              },
              outputs=output)
          model_with_dict_input.compile('sgd', 'mse')
          model_with_dict_input.set_weights(model_weights)
          model_with_dict_input_fit = model_with_dict_input.fit(
              dataset, steps_per_epoch=1, epochs=1).history
          self.assertAllClose(
              model_with_dict_input_fit,
              model_with_array_input_fit,
              atol=1e-4,
              rtol=1e-4)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=strategies_minus_tpu, mode=['graph', 'eager']) +
        tf.__internal__.test.combinations.combine(
            distribution=multi_worker_mirrored_strategies, mode=['eager']))
    def test_dataset_with_sample_weights(self, distribution):
      with self.cached_session(), distribution.scope():
        model = get_sample_weights_model()
        optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate=0.001)
        loss = 'mse'
        model.compile(
            optimizer,
            loss)
  
        inputs = np.array([[0], [1], [2], [3]], np.float32)
        targets = np.array([[2], [4], [6], [8]], np.float32)
        sample_weights = np.array([0.25, 0.5, 0.75, 1], np.float32)
        ds = tf.data.Dataset.from_tensor_slices(
            (inputs, targets, sample_weights)).batch(2)
        result = model.evaluate(ds, verbose=1)
  
        # The per sample loss is multiplied by the corresponding sample weight.
        # The average of these weighted losses is the return value of the
        # `evaluate` call. For example, in the test above the average weighted
        # loss is calculated in the following manner:
        # batch_1 = (((2-0)^2) * 0.25 + ((4-1)^2) * 0.5) / 2 = 5.5 / 2 = 2.75
        # batch_2 = (((6-2)^2 * 0.75) + ((8-3)^2 * 1)) / 2 = 37 / 2 = 18.5
        # final result = (batch_1 + batch_2) / 2 = 10.625.
        # The first time we divide by number of input samples and the second time
        # we divide by number of steps/batches that the loss is aggregated over.
        self.assertAllClose(result, 10.625)
  
        # We now test without passing sample_weights:
        # batch_1 = ((2-0)^2) + ((4-1)^2) / 2 = 13 / 2 = 6.5
        # batch_2 = ((6-2)^2) + ((8-3)^2) / 2 = 41 / 2 = 20.5
        # final result = (batch_1 + batch_2) / 2 =  27 / 2 = 13.5
        ds = tf.data.Dataset.from_tensor_slices((inputs, targets)).batch(2)
        result = model.evaluate(ds, verbose=1)
        self.assertAllClose(result, 13.5)
  
  
  class TestDistributionStrategyWithDatasetsFile(tf.test.TestCase,
                                                 parameterized.TestCase):
  
    def setUp(self):
      super(TestDistributionStrategyWithDatasetsFile, self).setUp()
      self.input_file_name = os.path.join(self.get_temp_dir(), 'input.tfrecord')
      inputs = np.zeros((20, 3), dtype=np.float32)
      input_dataset = tf.data.Dataset.from_tensor_slices(inputs)
      input_dataset = input_dataset.map(tf.io.serialize_tensor)
      writer = tf.data.experimental.TFRecordWriter(self.input_file_name)
      writer.write(input_dataset)
  
    # TODO(wxinyi): add a multi-worker test for TPU
    @tf.__internal__.distribute.combinations.generate(multi_worker_strategy_combinations_eager_only())
    def test_predict_on_dataset_shard_options_file_multi_worker_mirrored(
        self, distribution, mode):
      # This test is to verify if we successfully switch auto_shard_policy of a
      # input dataset inside model.predict with MultiWorkerMirroredStrategy to
      # AutoShardPolicy.DATA. Since there is only one input file for multiple
      # workers, AutoShardPolicy.AUTO or AutoShardPolicy.FILE will lead to an
      # error. However, since we switch to AutoShardPolicy.DATA in model.predict,
      # no error is raised.
      del mode
      with distribution.scope():
        optimizer_fn = gradient_descent_keras.SGD
        optimizer = optimizer_fn(0.001)
        model = get_model()
        loss = 'mse'
        model.compile(optimizer, loss)
  
      dataset = tf.data.TFRecordDataset(self.input_file_name)
      dataset = dataset.map(lambda x: tf.io.parse_tensor(x, tf.float32))
  
      dummy_op = lambda inp: True
  
      dataset = dataset.filter(dummy_op).batch(8, drop_remainder=True)
  
      options = tf.data.Options()
      options.experimental_distribute.auto_shard_policy = \
          tf.data.experimental.AutoShardPolicy.FILE
      dataset = dataset.with_options(options)
  
      model.predict(dataset, steps=1)
  
  
  class TestRegularizerLoss(tf.test.TestCase, parameterized.TestCase):
  
    class IdentityRegularizer(keras.regularizers.Regularizer):
  
      def __call__(self, x):
        return tf.identity(x)
  
    class AddLayer(keras.layers.Layer):
  
      def build(self, _):
        self.v = self.add_weight(
            'v', (),
            initializer='ones',
            regularizer=TestRegularizerLoss.IdentityRegularizer())
  
      def call(self, inputs):
        return inputs + self.v
  
    @staticmethod
    def loss_fn(_, y_pred):
      return tf.reduce_mean(y_pred)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(all_strategy_combinations_minus_default()))
    def test_regularizer_loss(self, distribution):
      batch_size = 2
      if not distributed_training_utils.global_batch_size_supported(distribution):
        batch_size //= distribution.num_replicas_in_sync
  
        # Given an input x, which is always 1, and variable v, this model computes
        # Loss=x+v+regularizer_loss, where regularizer_loss=v and the variable is
        # initialized to 1. Therefore, this model computes Loss=1+2v, and so the
        # gradient dLoss/dv = 2. This gradient of 2 is averaged over all examples
        # in a batch and then multiplied by the learning rate of 1. As a result,
        # the model update for one batch should subtract 2 from v, resulting in v
        # being -1. If the regularizer loss is not scaled correctly by number of
        # replicas, the variable value will be incorrect when number of replicas
        # >1. For e.g. it will be -2 if num replicas = 2.
      with distribution.scope():
        x = keras.layers.Input(shape=(1,), batch_size=batch_size)
        y = TestRegularizerLoss.AddLayer()(x)
        model = keras.models.Model(inputs=x, outputs=y)
        opt = gradient_descent_keras.SGD(1.)
        model.compile(
            opt,
            loss=TestRegularizerLoss.loss_fn)
        model.fit(
            x=np.array([[1.], [1.]], dtype=np.float32),
            y=np.array([[1.], [1.]], dtype=np.float32),
            batch_size=batch_size)
        v = model.get_weights()[0]
        self.assertEqual(-1.0, v)
  
  
  @testing_utils.run_all_without_tensor_float_32(
      'Uses Dense layers, which call matmul')
  class TestDistributionStrategyWithKerasModels(tf.test.TestCase,
                                                parameterized.TestCase):
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_distribution_strategy_on_sequential_model(
        self, distribution):
      with distribution.scope():
        optimizer_fn = gradient_descent_keras.SGD
        optimizer = optimizer_fn(learning_rate=0.001)
        model = simple_sequential_model()
        loss = 'mse'
        model.compile(
            optimizer,
            loss)
  
        inputs = np.zeros((20, 10), np.float32)
        targets = np.zeros((20, 2), np.float32)
  
      model.fit(inputs, targets, epochs=1, batch_size=10)
      model.predict(inputs, batch_size=10)
      model.evaluate(inputs, targets, batch_size=10)
  
    @tf.__internal__.distribute.combinations.generate(all_strategy_combinations())
    def test_distribution_strategy_on_functional_model(
        self, distribution):
      with distribution.scope():
        optimizer_fn = gradient_descent_keras.SGD
        optimizer = optimizer_fn(learning_rate=0.001)
        model = get_model()
        loss = 'mse'
        model.compile(
            optimizer,
            loss)
  
        inputs = np.zeros((64, 3), dtype=np.float32)
        targets = np.zeros((64, 4), dtype=np.float32)
  
      model.fit(inputs, targets, epochs=1)
      model.predict(inputs)
      model.evaluate(inputs, targets)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_distributed_dataset(self, distribution):
      with distribution.scope():
  
        class CBCounter(keras.callbacks.Callback):
  
          def __init__(self):
            self.epochs = 0
            self.train_batches = 0
            self.test_batches = 0
  
          def on_epoch_end(self, batch, logs=None):
            self.epochs += 1
  
          def on_train_batch_end(self, batch, logs=None):
            self.train_batches += 1
  
          def on_test_batch_end(self, batch, logs=None):
            self.test_batches += 1
  
        model = keras.Sequential([keras.layers.Dense(1)])
        model.compile('sgd', 'mse')
        cb_counter = CBCounter()
  
        x, y = np.ones((100, 10)), np.ones((100, 1))
        ds = tf.data.Dataset.from_tensor_slices((x, y))
        ds = ds.batch(10).repeat(2)
        ds = distribution.experimental_distribute_dataset(ds)
  
        val_ds = tf.data.Dataset.from_tensor_slices((x, y))
        val_ds = val_ds.batch(20)
        val_ds = distribution.experimental_distribute_dataset(val_ds)
  
        model.fit(
            ds,
            steps_per_epoch=10,
            validation_data=val_ds,
            validation_steps=5,
            epochs=2,
            callbacks=[cb_counter])
  
        self.assertEqual(cb_counter.train_batches, 20)
        self.assertEqual(cb_counter.test_batches, 10)
        self.assertEqual(cb_counter.epochs, 2)
  
        # Check for `steps_per_epoch`.
        if distribution.num_replicas_in_sync > 1:
          with self.assertRaisesRegex(ValueError,
                                      'distributed dataset, you must specify'):
            model.fit(ds, epochs=2)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_distributed_datasets_from_function(self, distribution):
      with distribution.scope():
  
        class CBCounter(keras.callbacks.Callback):
  
          def __init__(self):
            self.epochs = 0
            self.train_batches = 0
            self.test_batches = 0
  
          def on_epoch_end(self, batch, logs=None):
            self.epochs += 1
  
          def on_train_batch_end(self, batch, logs=None):
            self.train_batches += 1
  
          def on_test_batch_end(self, batch, logs=None):
            self.test_batches += 1
  
        model = keras.Sequential([keras.layers.Dense(1)])
        model.compile('sgd', 'mse')
        cb_counter = CBCounter()
  
        def make_dataset(_):
          x, y = np.ones((100, 10)), np.ones((100, 1))
          ds = tf.data.Dataset.from_tensor_slices((x, y))
          ds = ds.batch(5).repeat()
          return ds
  
        ds = distribution.distribute_datasets_from_function(make_dataset)
        val_ds = distribution.distribute_datasets_from_function(make_dataset)
  
        model.fit(
            ds,
            steps_per_epoch=10,
            validation_data=val_ds,
            validation_steps=5,
            epochs=2,
            callbacks=[cb_counter])
  
        self.assertEqual(cb_counter.train_batches, 20)
        self.assertEqual(cb_counter.test_batches, 10)
        self.assertEqual(cb_counter.epochs, 2)
  
        # Check for `steps_per_epoch`.
        if distribution.num_replicas_in_sync > 1:
          with self.assertRaisesRegex(ValueError,
                                      'distributed dataset, you must specify'):
            model.fit(ds, epochs=2)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_host_training_loop(self, distribution):
      if isinstance(distribution,
                    tf.distribute.MultiWorkerMirroredStrategy):
        self.skipTest('b/172032817')
      with distribution.scope():
        inputs = keras.Input((10, 10, 3))
        x = keras.layers.Conv2D(3, kernel_size=3)(inputs)
        x = keras.layers.Flatten()(x)
        outputs = keras.layers.Dense(1)(x)
        model = keras.Model(inputs, outputs)
  
      model.compile('sgd', 'mse', steps_per_execution=10)
  
      bc = BatchCountingCB()
      x, y = np.ones((100, 10, 10, 3)), np.ones((100, 1))
      model.fit(x, y, batch_size=2, epochs=1, callbacks=[bc])
      self.assertEqual(bc.train_begin_batches, [0, 10, 20, 30, 40])
      self.assertEqual(bc.train_end_batches, [9, 19, 29, 39, 49])
  
      model.evaluate(x, y, batch_size=2, callbacks=[bc])
      self.assertEqual(bc.test_begin_batches, [0, 10, 20, 30, 40])
      self.assertEqual(bc.test_end_batches, [9, 19, 29, 39, 49])
  
      model.predict(x, batch_size=2, callbacks=[bc])
      self.assertEqual(bc.predict_begin_batches, [0, 10, 20, 30, 40])
      self.assertEqual(bc.predict_end_batches, [9, 19, 29, 39, 49])
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_host_training_loop_last_partial_execution(self, distribution):
      if isinstance(distribution,
                    tf.distribute.MultiWorkerMirroredStrategy):
        self.skipTest('b/172032817')
      with distribution.scope():
        inputs = keras.Input(10)
        outputs = keras.layers.Dense(1)(inputs)
        model = keras.Model(inputs, outputs)
  
      model.compile('sgd', 'mse', steps_per_execution=20)
  
      bc = BatchCountingCB()
      x, y = np.ones((100, 10)), np.ones((100, 1))
      model.fit(x, y, batch_size=2, epochs=1, callbacks=[bc])
      self.assertEqual(bc.train_begin_batches, [0, 20, 40])
      self.assertEqual(bc.train_end_batches, [19, 39, 49])
  
      model.evaluate(x, y, batch_size=2, callbacks=[bc])
      self.assertEqual(bc.test_begin_batches, [0, 20, 40])
      self.assertEqual(bc.test_end_batches, [19, 39, 49])
  
      model.predict(x, batch_size=2, callbacks=[bc])
      self.assertEqual(bc.predict_begin_batches, [0, 20, 40])
      self.assertEqual(bc.predict_end_batches, [19, 39, 49])
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_host_training_loop_dataset_unknown_size(self, distribution):
      if isinstance(distribution,
                    tf.distribute.MultiWorkerMirroredStrategy):
        self.skipTest('b/172032817')
      with distribution.scope():
        inputs = keras.Input(10)
        outputs = keras.layers.Dense(1)(inputs)
        model = keras.Model(inputs, outputs)
  
      model.compile('sgd', 'mse', steps_per_execution=20)
  
      x, y = np.ones((100, 10)), np.ones((100, 1))
      ds = tf.data.Dataset.from_tensor_slices((x, y)).batch(2)
      ds = ds.filter(lambda *args, **kwargs: True)  # Makes the size UNKNOWN.
      bc = BatchCountingCB()
  
      with self.assertRaisesRegex(ValueError, 'steps_per_execution'):
        model.fit(ds, epochs=2, callbacks=[bc])
  
      train_ds = ds.repeat(2)
      model.fit(train_ds, steps_per_epoch=50, epochs=2, callbacks=[bc])
      self.assertEqual(bc.train_begin_batches, [0, 20, 40, 0, 20, 40])
      self.assertEqual(bc.train_end_batches, [19, 39, 49, 19, 39, 49])
  
      with self.assertRaisesRegex(ValueError, 'steps_per_execution'):
        model.evaluate(ds, callbacks=[bc])
  
      test_ds = ds.repeat(2)
      model.evaluate(test_ds, steps=50, callbacks=[bc])
      self.assertEqual(bc.test_begin_batches, [0, 20, 40])
      self.assertEqual(bc.test_end_batches, [19, 39, 49])
  
      predict_ds = ds.repeat(2)
      model.predict(predict_ds, steps=50, callbacks=[bc])
      self.assertEqual(bc.predict_begin_batches, [0, 20, 40])
      self.assertEqual(bc.predict_end_batches, [19, 39, 49])
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_host_training_loop_truncate_to_epoch(self, distribution):
      if isinstance(distribution,
                    tf.distribute.MultiWorkerMirroredStrategy):
        self.skipTest('b/172032817')
      with distribution.scope():
        inputs = keras.Input(10)
        outputs = keras.layers.Dense(1)(inputs)
        model = keras.Model(inputs, outputs)
  
      model.compile('sgd', 'mse', steps_per_execution=500)
  
      x, y = np.ones((100, 10)), np.ones((100, 1))
      bc = BatchCountingCB()
      model.fit(x, y, batch_size=2, epochs=2, callbacks=[bc])
      self.assertEqual(bc.train_begin_batches, [0, 0])
      self.assertEqual(bc.train_end_batches, [49, 49])
  
      x, y = np.ones((50, 10)), np.ones((50, 1))
      model.evaluate(x, y, batch_size=2, callbacks=[bc])
      self.assertEqual(bc.test_begin_batches, [0])
      self.assertEqual(bc.test_end_batches, [24])
  
      x = np.ones((50, 10))
      model.predict(x, batch_size=2, callbacks=[bc])
      self.assertEqual(bc.predict_begin_batches, [0])
      self.assertEqual(bc.predict_end_batches, [24])
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_gradient_clipping(self, distribution):
  
      class MyLayer(keras.layers.Layer):
  
        def build(self, _):
          self.v1 = tf.Variable(1.)
          self.v2 = tf.Variable(1.)
  
        def call(self, x):
          return 3 * self.v1 - 3 * self.v2
  
      x, y = np.ones((10, 1)), np.ones((10, 1))
  
      with distribution.scope():
        layer = MyLayer()
        model = keras.Sequential([layer])
        optimizer = gradient_descent_keras.SGD(1., clipnorm=2., clipvalue=2.)
      model.compile(optimizer, 'mae')
  
      if isinstance(distribution,
                    (tf.distribute.experimental.CentralStorageStrategy,
                     tf.compat.v1.distribute.experimental.CentralStorageStrategy)):
        with self.assertRaisesRegex(ValueError, 'not supported'):
          model.fit(x, y, batch_size=10, epochs=1)
      else:
        model.fit(x, y, batch_size=10, epochs=1)
        self.assertAllClose(self.evaluate(layer.v1), 3.)
        self.assertAllClose(self.evaluate(layer.v2), -1.)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_custom_gradient_transformation(self, distribution):
      if isinstance(distribution,
                    (tf.distribute.experimental.CentralStorageStrategy,
                     tf.compat.v1.distribute.experimental.CentralStorageStrategy)):
        self.skipTest('Not supported with `CentralStorageStrategy`')
  
      class MyLayer(keras.layers.Layer):
  
        def build(self, _):
          self.v1 = tf.Variable(1.)
          self.v2 = tf.Variable(-1.)
  
        def call(self, x):
          return x + self.v1 + self.v2
  
      def custom_transform(grads_and_vars):
        # Always set gradients to 1.
        return [(tf.ones_like(g), v) for g, v in grads_and_vars]
  
      x, y = np.ones((10, 1)), np.ones((10, 1))
  
      with distribution.scope():
        layer = MyLayer()
        model = keras.Sequential([layer])
        optimizer = gradient_descent_keras.SGD(
            1., gradient_transformers=[custom_transform])
      model.compile(optimizer, 'mae')
  
      model.fit(x, y, batch_size=10, epochs=1)
      self.assertAllClose(self.evaluate(layer.v1), 0.)
      self.assertAllClose(self.evaluate(layer.v2), -2.)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(
            all_strategy_combinations_minus_default()))
    def test_distribution_strategy_one_dimensional(self, distribution):
      with distribution.scope():
        inp = keras.layers.Input(shape=(10,))
        out = keras.layers.Dense(3, activation='softmax')(inp)
        model = keras.Model(inputs=[inp], outputs=[out])
        model.compile(
            optimizer='rmsprop',
            loss='sparse_categorical_crossentropy',
            metrics=['sparse_categorical_accuracy'])
  
        x = np.random.random((64, 10)).astype('float32')
        y = np.random.randint(3, size=64)
  
        model.fit(x, y, epochs=1, steps_per_epoch=2)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=[
                tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
                tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus
            ],
            mode=['graph', 'eager'],
            reduction=[
                losses_utils.ReductionV2.AUTO,
                losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE,
                losses_utils.ReductionV2.SUM
            ]))
    def test_distribution_strategy_with_loss_reduction_types(
        self, distribution, reduction):
      np.random.seed(_RANDOM_SEED)
  
      def _get_model():
        inputs = keras.Input((10,))
        x1 = keras.layers.Dense(10, kernel_initializer='zeros')(inputs)
        x2 = keras.layers.Dense(10, kernel_initializer='zeros')(x1)
        outputs = keras.layers.Dense(1, kernel_initializer='zeros')(x2)
        model = keras.Model(inputs, outputs)
        return model
  
      x = np.random.random((64, 10))
      y = np.random.random((64, 1))
      dataset = tf.data.Dataset.from_tensor_slices((x, y))
      dataset = dataset.batch(32)
  
      model = _get_model()
      model.compile(
          'sgd', loss=keras.losses.MeanSquaredError(reduction=reduction))
      history = model.fit(dataset, steps_per_epoch=2, epochs=1, shuffle=False)
  
      with distribution.scope():
        ds_model = _get_model()
        ds_model.compile(
            'sgd',
            loss=keras.losses.MeanSquaredError(reduction=reduction))
        ds_history = ds_model.fit(
            dataset, steps_per_epoch=2, epochs=1, shuffle=False)
      self.assertArrayNear(history.history['loss'], ds_history.history['loss'],
                           1e-5)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(
            all_strategy_combinations_minus_default()))
    def test_distribution_strategy_with_symbolic_add_loss(
        self, mode, distribution):
  
      def _make_model_with_add_loss():
        inputs = keras.Input((10,))
        x1 = keras.layers.Dense(10, kernel_initializer='zeros')(inputs)
        x2 = keras.layers.Dense(10, kernel_initializer='zeros')(x1)
        outputs = keras.layers.Dense(1, kernel_initializer='zeros')(x2)
        model = keras.Model(inputs, outputs)
        model.add_loss(tf.reduce_mean(x1))
        model.add_loss(tf.reduce_mean(outputs))
        return model
  
      x = np.ones((64, 10)).astype('float32')
  
      model = _make_model_with_add_loss()
      model.compile('sgd')
      history = model.fit(x, epochs=1)
  
      with distribution.scope():
        ds_model = _make_model_with_add_loss()
        ds_model.compile(
            'sgd')
        ds_history = ds_model.fit(x, epochs=1)
  
      self.assertAllClose(history.history, ds_history.history)
  
    # TODO(omalleyt): Investigate flakiness and re-enable.
    @tf.__internal__.distribute.combinations.generate(all_strategy_minus_default_and_tpu_combinations())
    def DISABLED_test_distribution_strategy_with_callable_add_loss(
        self, distribution):
  
      def _make_model():
        inputs = keras.Input((10,))
        x1 = keras.layers.Dense(10, kernel_initializer='zeros')(inputs)
        x2 = keras.layers.Dense(10, kernel_initializer='zeros')(x1)
        d = keras.layers.Dense(1, kernel_initializer='zeros')
        outputs = d(x2)
        model = keras.Model(inputs, outputs)
        model.add_loss(lambda: 100. * tf.reduce_mean(d.kernel))
        return model
  
      x = np.ones((64, 10)).astype('float32')
      y = np.ones((64, 1)).astype('float32')
  
      model = _make_model()
      self.assertLen(model.losses, 1)
  
      model.compile('sgd', 'mse')
      history = model.fit(x, y, steps_per_epoch=2, epochs=1)
  
      with distribution.scope():
        ds_model = _make_model()
        self.assertLen(ds_model.losses, 1)
        ds_model.compile('sgd', 'mse')
        ds_history = ds_model.fit(x, y, steps_per_epoch=2, epochs=1)
  
      self.assertAllClose(history.history, ds_history.history)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(
            all_strategy_minus_default_and_tpu_combinations()))
    def test_distribution_strategy_with_add_metric_in_call(
        self, distribution):
  
      class Bias(keras.layers.Layer):
  
        def build(self, input_shape):
          self.bias = self.add_weight(name='bias', initializer='zeros', shape=())
  
        def call(self, inputs):
          self.add_metric(
              tf.reduce_mean(inputs), name='bias', aggregation='mean')
          return inputs + self.bias
  
      def _make_model_with_add_metric():
        inputs = keras.Input((10,))
        x1 = keras.layers.Dense(10, kernel_initializer='zeros')(inputs)
        x2 = Bias()(x1)
        outputs = keras.layers.Dense(1, kernel_initializer='zeros')(x2)
        model = keras.Model(inputs, outputs)
        return model
  
      x = np.ones((64, 10)).astype('float32')
      y = np.ones((64, 1)).astype('float32')
  
      model = _make_model_with_add_metric()
      self.assertLen(model.metrics, 1)
  
      model.compile('sgd', 'mse')
      history = model.fit(
          x, y, validation_data=(x, y), validation_steps=2, epochs=2)
  
      with distribution.scope():
        ds_model = _make_model_with_add_metric()
        self.assertLen(ds_model.metrics, 1)
        ds_model.compile(
            'sgd',
            'mse')
        ds_history = ds_model.fit(
            x, y, validation_data=(x, y), validation_steps=2, epochs=2)
        # includes stateful loss metric in eager.
        metrics_len = 2 if tf.executing_eagerly() else 1
        self.assertLen(ds_model.metrics, metrics_len)
  
      self.assertAllClose(history.history, ds_history.history)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=[
                tf.__internal__.distribute.combinations.one_device_strategy,
                tf.__internal__.distribute.combinations.one_device_strategy_gpu,
                tf.__internal__.distribute.combinations.mirrored_strategy_with_gpu_and_cpu,
                tf.__internal__.distribute.combinations.mirrored_strategy_with_two_gpus
            ],
            mode=['eager']))
    def test_distribution_strategy_with_add_metric_object(
        self, distribution):
  
      class Bias(keras.layers.Layer):
  
        def build(self, input_shape):
          self.bias = self.add_weight(name='bias', initializer='zeros', shape=())
          self.mean = keras.metrics.Mean(name='mean')
  
        def call(self, inputs):
          self.add_metric(self.mean(inputs))
          return inputs + self.bias
  
      def _make_model_with_add_metric_object():
        inputs = keras.Input((10,))
        x1 = keras.layers.Dense(10, kernel_initializer='zeros')(inputs)
        x2 = Bias()(x1)
        outputs = keras.layers.Dense(1, kernel_initializer='zeros')(x2)
        model = keras.Model(inputs, outputs)
        return model
  
      x = np.ones((64, 10)).astype('float32')
      y = np.ones((64, 1)).astype('float32')
  
      model = _make_model_with_add_metric_object()
      self.assertLen(model.metrics, 1)
  
      model.compile('sgd', 'mse')
      history = model.fit(
          x, y, validation_data=(x, y), validation_steps=2, epochs=2)
  
      with distribution.scope():
        ds_model = _make_model_with_add_metric_object()
        self.assertLen(ds_model.metrics, 1)
        ds_model.compile(
            'sgd',
            'mse')
        ds_history = ds_model.fit(
            x, y, validation_data=(x, y), validation_steps=2, epochs=2)
        # includes stateful loss metric in eager.
        metrics_len = 2 if tf.executing_eagerly() else 1
        self.assertLen(ds_model.metrics, metrics_len)
  
      self.assertAllClose(history.history, ds_history.history)
  
    @tf.__internal__.distribute.combinations.generate(
        # TODO(phillypham): Why does validation_steps > 1 not work on TPUs?
        tf.__internal__.test.combinations.times(
            all_strategy_minus_default_and_tpu_combinations()))
    def test_distribution_strategy_with_add_metric_outside_call(
        self, distribution):
  
      def _make_model_with_add_metric():
        inputs = keras.Input((10,))
        x1 = keras.layers.Dense(10, kernel_initializer='zeros')(inputs)
        outputs = keras.layers.Dense(1, kernel_initializer='zeros')(x1)
        model = keras.Model(inputs, outputs)
        model.add_metric(
            tf.reduce_mean(x1), name='mid_mean', aggregation='mean')
        return model
  
      x = np.ones((64, 10)).astype('float32')
      y = np.ones((64, 1)).astype('float32')
  
      model = _make_model_with_add_metric()
      self.assertLen(model.metrics, 1)
  
      model.compile('sgd', 'mse')
      history = model.fit(
          x, y, validation_data=(x, y), validation_steps=2, epochs=2)
  
      with distribution.scope():
        ds_model = _make_model_with_add_metric()
        self.assertLen(ds_model.metrics, 1)
        ds_model.compile(
            'sgd',
            'mse')
        ds_history = ds_model.fit(
            x, y, validation_data=(x, y), validation_steps=2, epochs=2)
        # includes stateful loss metric in eager.
        metrics_len = 2 if tf.executing_eagerly() else 1
        self.assertLen(ds_model.metrics, metrics_len)
  
      self.assertAllClose(history.history, ds_history.history)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=strategies_minus_tpu + multi_worker_mirrored_strategies,
            mode=['eager']))
    def test_sparse_tensor_outputs(self, distribution):
  
      class ToSparse(keras.layers.Layer):
        """Create a sparse tensor based on a given dense tensor."""
  
        def call(self, inputs):
          indices = tf.where(tf.not_equal(inputs, 0))
          values = tf.gather_nd(inputs, indices)
          shape = tf.shape(inputs, out_type='int64')
          return tf.SparseTensor(indices, values, dense_shape=shape)
  
      model = keras.Sequential([ToSparse()])
  
      # Define some input data with additional padding.
      input_data = np.array([[1, 0, 0], [2, 3, 0]])
      output = model.predict(input_data, batch_size=2)
  
      expected_indices = np.array([[0, 0], [1, 0], [1, 1]])
      expected_values = np.array([1, 2, 3])
      expected_dense_shape = np.array([2, 3])
  
      self.assertAllEqual(output.indices, expected_indices)
      self.assertAllEqual(output.values, expected_values)
      self.assertAllEqual(output.dense_shape, expected_dense_shape)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=strategies_minus_tpu + multi_worker_mirrored_strategies,
            mode=['eager']))
    def test_ragged_tensor_outputs(self, distribution):
  
      class ToRagged(keras.layers.Layer):
        """Create a ragged tensor based on a given dense tensor."""
  
        def __init__(self, padding, ragged_rank=1, **kwargs):
          super(ToRagged, self).__init__(**kwargs)
          self._padding = padding
          self._ragged_rank = ragged_rank
  
        def call(self, inputs):
          return tf.RaggedTensor.from_tensor(
              inputs, padding=self._padding, ragged_rank=self._ragged_rank)
  
      model = keras.Sequential([ToRagged(padding=0)])
  
      # Define some input data with additional padding.
      input_data = np.array([[1, 0, 0], [2, 3, 0]])
      output = model.predict(input_data, batch_size=2)
  
      expected_values = [[1], [2, 3]]
      self.assertAllEqual(expected_values, output)
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=strategies_minus_default_minus_tpu + tpu_strategies +
            multi_worker_mirrored_strategies,
            mode=['eager']))
    def test_correctness_of_add_loss_with_merge_call(self, distribution):
      batch_size = 32
  
      def _get_model():
        inputs = keras.layers.Input(shape=(1,))
        labels = keras.layers.Input(shape=(1,))
        x = keras.layers.Dense(10, activation='relu')(inputs)
        y = keras.layers.Dense(1)(x)
        model = keras.models.Model([inputs, labels], y)
        model.add_loss(keras.losses.mean_squared_error(labels, y))
        return model
  
      def _get_data():
        x_train = np.random.rand(64, 1)
        y_train = 3 * x_train
        x_train = x_train.astype('float32')
        y_train = y_train.astype('float32')
        dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))
        dataset = dataset.batch(batch_size)
        return dataset
  
      with distribution.scope():
        model = _get_model()
        optimizer = gradient_descent_keras.SGD(0.2)
  
        @tf.function
        def train_step(dist_inputs):
  
          def step_fn(inputs):
            with tf.GradientTape() as tape:
              logits = model(inputs)
  
              # Invoke a merge_call()
              tf.distribute.get_replica_context().merge_call(
                  lambda d: None)
  
              # Verify that there is only one loss on the model.
              assert len(model.losses) == 1
              loss_from_model = tf.reduce_sum(
                  model.losses) * 1.0 / batch_size
  
              # Compute loss in this loop.
              loss = keras.losses.mean_squared_error(inputs[1], logits)
              loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)
  
              # Verify that the loss computed in this loop is equivalent to the
              # loss from the model that was added via add_loss.
              tf.compat.v1.assert_equal(loss, loss_from_model)
  
            grads = tape.gradient(loss, model.trainable_variables)
            optimizer.apply_gradients(zip(grads, model.trainable_variables))
            return loss
  
          per_replica_losses = distribution.run(step_fn, args=(dist_inputs,))
          return distribution.reduce(
              tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)
  
        dataset = distribution.experimental_distribute_dataset(_get_data())
        for _ in range(2):
          for x in dataset:
            train_step(x)
  
    @tf.__internal__.distribute.combinations.generate(tf.__internal__.test.combinations.combine(mode=['graph', 'eager']))
    def test_unimplemented_parameter_server_strategy(self):
      cluster_spec = multi_worker_testing_utils.create_in_process_cluster(
          num_workers=3, num_ps=2)
      cluster_resolver = SimpleClusterResolver(
          cluster_spec=tf.train.ClusterSpec(cluster_spec),
          task_type='worker',
          task_id=1,
          num_accelerators={'GPU': 0})
      distribution = tf.compat.v1.distribute.experimental.ParameterServerStrategy(
          cluster_resolver)
  
      self.assertIsInstance(distribution,
                            tf.compat.v1.distribute.experimental.ParameterServerStrategy)
  
      with self.assertRaisesRegex(NotImplementedError,
                                  'ParameterServerStrategy*'):
        with distribution.scope():
          model = simple_sequential_model()
          optimizer = tf.compat.v1.train.RMSPropOptimizer(learning_rate=0.001)
          loss = 'mse'
          model.compile(optimizer, loss)
  
  
  # Models to exercise inserting ancillary layers with add_loss and add_metric.
  def _functional_with_add_loss_and_metric(input_shape, num_classes, l1, l2):
    inputs = keras.Input(input_shape, name='images')
    x = keras.layers.Conv2D(32, kernel_size=5, activation='relu')(inputs)
    x = keras.layers.MaxPooling2D(pool_size=2)(x)
    x = keras.layers.Conv2D(64, kernel_size=5, activation='relu')(x)
    x = keras.layers.MaxPooling2D(pool_size=2)(x)
    # Apply L2 regularization to embedding. Use a mix of TensorFlow ops and layers
    # to exercise all code paths.
    x = keras.layers.Flatten(name='embedding')(x)
    l2_loss = tf.reduce_mean(tf.reduce_sum(tf.square(x), -1))
    # Apply L1 regularization to next layer.
    x = keras.layers.Dense(1024, activation='relu', name='sparse_embedding')(x)
    l1_loss = keras.layers.Lambda(
        lambda x: tf.reduce_mean(tf.reduce_sum(x, -1)),
        name='l1_loss')(
            x)
    outputs = keras.layers.Dense(num_classes, name='logits')(x)
    model = keras.Model(inputs=inputs, outputs=outputs)
    # Weight regularization terms.
    model.add_loss(keras.layers.Lambda(lambda x: x * l2)(l2_loss))
    model.add_metric(l2_loss, aggregation='mean', name='l2_loss')
    model.add_loss(l1_loss * l1)
    model.add_metric(l1_loss, aggregation='mean', name='l1_loss')
    return model
  
  
  def _sequential_with_add_loss_and_metric(input_shape, num_classes, l1, l2):
    model = keras.Sequential([
        keras.layers.Conv2D(
            32, kernel_size=5, activation='relu', input_shape=input_shape),
        keras.layers.MaxPooling2D(pool_size=2),
        keras.layers.Conv2D(64, kernel_size=5, activation='relu'),
        keras.layers.MaxPooling2D(pool_size=2),
        keras.layers.Flatten(name='embedding'),
        keras.layers.Dense(1024, activation='relu', name='sparse_embedding'),
        keras.layers.Dense(num_classes, name='logits'),
    ])
    # Extract layer outputs, add regularization terms, and rescale the metric.
    # Use a mix of TensorFlow ops and layers to exercise all code paths.
    x = model.get_layer('sparse_embedding').get_output_at(-1)
    l1_loss = l1 * tf.reduce_mean(tf.reduce_sum(x, -1))
    model.add_loss(l1_loss)
    model.add_metric(
        keras.layers.Lambda(lambda x: tf.divide(x, l1))(l1_loss),
        aggregation='mean',
        name='l1_loss')
    x = model.get_layer('embedding').get_output_at(-1)
    l2_loss = keras.layers.Lambda(
        lambda x: l2 * tf.reduce_mean(tf.reduce_sum(x * x, -1)),
        name='l2_loss')(
            x)
    model.add_loss(l2_loss)
    model.add_metric(l2_loss / l2, aggregation='mean', name='l2_loss')
    return model
  
  
  def _functional_with_layer_reuse(input_shape, num_classes, l1, l2):
    base_model = keras.Sequential([
        keras.layers.Conv2D(
            32, kernel_size=5, activation='relu', input_shape=input_shape),
        keras.layers.MaxPooling2D(pool_size=2),
        keras.layers.Conv2D(64, kernel_size=5, activation='relu'),
        keras.layers.MaxPooling2D(pool_size=2),
        keras.layers.Flatten(),
        keras.layers.Dense(1024, activation='relu'),
        keras.layers.Dense(num_classes, name='logits'),
    ])
    inputs = keras.Input(input_shape, name='images')
    logits = base_model(inputs)
    model = keras.Model(inputs=inputs, outputs=logits)
    # Reuse sequential layer and create new nodes.
    zero_logits = base_model(tf.zeros_like(inputs))
    one_logits = base_model(tf.ones_like(inputs))
    # L2 loss.
    l2_loss = tf.reduce_mean(
        tf.reduce_sum(tf.square(logits - zero_logits), -1))
    model.add_loss(l2_loss * l2)
    model.add_metric(l2_loss, aggregation='mean', name='l2_loss')
    # L1 loss.
    l1_loss = tf.reduce_mean(
        tf.reduce_sum(tf.abs(logits - one_logits), -1))
    model.add_loss(l1_loss * l1)
    model.add_metric(l1_loss, aggregation='mean', name='l1_loss')
    return model
  
  
  class TestDistributionStrategyWithMultipleAddLossAndMetricCalls(
      tf.test.TestCase, parameterized.TestCase):
    """Tests complex models with multiple add loss and metric calls."""
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.times(
            all_strategy_combinations_minus_default(),
            tf.__internal__.test.combinations.combine(
                model_fn=[
                    _functional_with_add_loss_and_metric,
                    _sequential_with_add_loss_and_metric,
                    _functional_with_layer_reuse,
                ],
                l1=[0.01],
                l2=[0.1])))
    def test_fit_and_evaluate(self, distribution, model_fn, l1, l2):
      # Make fake MNIST-like image data.
      np.random.seed(_RANDOM_SEED)
      dataset = tf.data.Dataset.from_tensor_slices(
          (np.random.uniform(size=(64, 28, 28, 1)).astype(np.float32),
           np.random.randint(0, 10, size=(64,))))
      dataset = dataset.shuffle(64).batch(
          8 * distribution.num_replicas_in_sync, drop_remainder=True)
      # Make model with distribution strategy and initialize with dataset shape.
      input_shape = tf.data.experimental.get_structure(dataset)[0].shape[1:]
      with distribution.scope():
        model = model_fn(input_shape, 10, l1, l2)
        model.compile(
            optimizer=keras.optimizers.adam_v2.Adam(1e-4),
            loss=keras.losses.SparseCategoricalCrossentropy(
                from_logits=True,
                reduction=losses_utils.ReductionV2.SUM_OVER_BATCH_SIZE),
            metrics=[
                keras.metrics.SparseCategoricalAccuracy(),
                keras.metrics.SparseCategoricalCrossentropy(from_logits=True),
            ])
      # Non-eager training doesn't support steps_per_epoch=None.
      for unused_epoch in range(2):
        model.fit(dataset)
      results = dict(zip(model.metrics_names, model.evaluate(dataset)))
      # Sanity checks.
      self.assertBetween(results['sparse_categorical_accuracy'], 0.02, 1.)
      self.assertGreater(results['l2_loss'], 0.)
      self.assertGreater(results['l1_loss'], 0.)
      # Assert correctness of the loss calculation and updating of metrics.
      self.assertNear(
          results['l1_loss'] * l1 + results['l2_loss'] * l2 +
          results['sparse_categorical_crossentropy'], results['loss'], 1e-6)
  
  
  class DeterministicModel(keras.Model):
    """Deterministic Model that always outputs the same initial result.
  
    It verifies the `call` method is run inside the same distribution
    strategy that the model was initially passed.
    """
  
    def __init__(self, strategy):
      super(DeterministicModel, self).__init__()
      self.x = None
      self.strategy = strategy
  
    def build(self, input_shape):
      self.x = tf.Variable(tf.ones(shape=()))
  
    def call(self, inputs, training=None, mask=None):
      active_strategy = tf.distribute.get_strategy()
      if active_strategy is not self.strategy:
        raise ValueError('Model must execute call w/ the original strategy')
      return self.x * inputs
  
  
  class TestModelCapturesStrategy(tf.test.TestCase, parameterized.TestCase):
    """Tests that model creation captures the strategy."""
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(distribution=all_strategies, mode=['eager']))
    def test_fit_and_evaluate(self, distribution):
      dataset = tf.data.Dataset.from_tensor_slices(
          (tf.ones(shape=(64,)), tf.ones(shape=(64,))))
      dataset = dataset.batch(8 * distribution.num_replicas_in_sync)
      # Make model with distribution strategy
      with distribution.scope():
        model = DeterministicModel(distribution)
        optimizer = keras.optimizers.adam_v2.Adam(1e-4)
  
      # Compile & evaluate the model outside of the distribution strategy scope
      model.compile(
          optimizer=optimizer,
          loss=keras.losses.MeanSquaredError(),
          metrics=['binary_accuracy'])
  
      # Call `optimizer.iterations` out of strategy scope.
      self.assertEqual(model.optimizer.iterations.numpy(), 0)
  
      # Non-eager training doesn't support steps_per_epoch=None.
      for unused_epoch in range(2):
        model.fit(dataset)
  
      results = model.evaluate(dataset)
      results = dict(zip(model.metrics_names, results))
  
      # Check that the metrics have a result we expect
      self.assertEqual(results['binary_accuracy'], 1.0)
      self.assertAllClose(results['loss'], 0.0)
  
      # Assert that all metric/optimizer/model variables were made in the
      # distribution strategy (Test that compile uses the captured
      # distribution strategy)
      metric_vars = tf.nest.flatten(
          [metric.variables for metric in model.metrics])
      for var in metric_vars:
        self.assertTrue(distribution.extended.variable_created_in_scope(var))
      for var in model.optimizer._weights:
        self.assertTrue(distribution.extended.variable_created_in_scope(var))
      for var in model.variables:
        self.assertTrue(distribution.extended.variable_created_in_scope(var))
  
      # Make sure the metric must be created in the same scope as the model:
      # This shouldn't raise any validation errors
      with distribution.scope():
        metric = keras.metrics.BinaryAccuracy()
      model.compile(
          optimizer=optimizer,
          loss=keras.losses.MeanSquaredError(),
          metrics=[metric])
  
      # This should raise an error because the metric is constructed
      # outside of the scope, and not by compile
      if tf.distribute.has_strategy():
        with self.assertRaisesRegex(ValueError, 'All metrics must be created in'):
          model.compile(
              optimizer=keras.optimizers.adam_v2.Adam(1e-4),
              loss=keras.losses.MeanSquaredError(),
              metrics=[keras.metrics.BinaryAccuracy()])
  
    @tf.__internal__.distribute.combinations.generate(
        tf.__internal__.test.combinations.combine(
            distribution=tf.__internal__.distribute.combinations.mirrored_strategy_with_one_cpu,
            mode=['eager']))
    def test_optimizer(self, distribution):
      temp_dir = os.path.join(self.get_temp_dir(), 'ckpt')
  
      def create_model():
        model = keras.models.Sequential([
            keras.layers.Dense(1),
        ])
        model.compile(optimizer='adam', loss='mse')
        model.build([None, 1])  # create weights.
        self.assertEmpty(model.optimizer.weights)
        return model
  
      model = create_model()
      x = y = tf.ones(shape=(1, 1))
      model.fit(x=x, y=y, batch_size=1)
      model.save_weights(temp_dir)
  
      with distribution.scope():
        model = create_model()
        model.load_weights(temp_dir)
        self.assertNotEmpty(model.optimizer.weights)
        self.assertTrue(
            distributed_training_utils.is_distributed_variable(
                model.optimizer.weights[0]))
  
      with distribution.scope():
        model = create_model()
      # create/restore slot variables outside of scope is fine.
      model.load_weights(temp_dir)
      self.assertNotEmpty(model.optimizer.weights)
      self.assertTrue(
          distributed_training_utils.is_distributed_variable(
              model.optimizer.weights[0]))
  
  
  if __name__ == '__main__':
    base_layer_utils.enable_v2_dtype_behavior()
    tf.__internal__.distribute.multi_process_runner.test_main()
  