@combinations.generate(strategy_and_inputs())
def test_correctness(self, distribution, use_numpy):
    with self.cached_session():
        tolerance = 1e-05
        if isinstance(distribution, mirrored_strategy.MirroredStrategy):
            tolerance = 0.0001
        keras.backend.set_image_data_format('channels_last')
        np.random.seed(_RANDOM_SEED)
        random_seed.set_random_seed(_RANDOM_SEED)
        num_samples = 9984
        x_train = np.random.rand(num_samples, 1)
        y_train = (3 * x_train)
        x_train = x_train.astype('float32')
        y_train = y_train.astype('float32')
        x_predict = [[1.0], [2.0], [3.0], [4.0]]
        model = keras.Sequential()
        model.add(keras.layers.Dense(10, activation='relu', input_shape=(1,)))
        model.add(keras.layers.Dense(10, activation='relu'))
        model.add(keras.layers.Dense(10, activation='relu'))
        model.add(keras.layers.Dense(1))
        initial_weights = model.get_weights()

        def fit_and_predict(with_distribution=None):
            model.set_weights(initial_weights)
            model.compile(loss=keras.losses.mean_squared_error, optimizer=gradient_descent.GradientDescentOptimizer(0.5), distribute=with_distribution)
            (training_inputs, eval_inputs, predict_inputs) = get_correctness_test_inputs(use_numpy, with_distribution, x_train, y_train, x_predict)
            model.fit(**training_inputs)
            eval_result = model.evaluate(**eval_inputs)
            weights = model.get_weights()
            predict_result = model.predict(**predict_inputs)
            return (weights, eval_result, predict_result)
        (wts_with_ds, eval_with_ds, predict_with_ds) = fit_and_predict(with_distribution=distribution)
        (wts_without_ds, eval_without_ds, predict_without_ds) = fit_and_predict(with_distribution=None)
        self.assertAllClose(wts_with_ds, wts_without_ds, atol=tolerance, rtol=tolerance)
        self.assertAllClose(eval_with_ds, eval_without_ds, atol=tolerance, rtol=tolerance)
        self.assertAllClose(predict_with_ds, predict_without_ds, atol=tolerance, rtol=tolerance)
@combinations.generate(strategy_and_inputs())
def test_correctness(self, distribution, use_numpy):
    with self.cached_session():
        tolerance = 1e-05
        if isinstance(distribution, mirrored_strategy.MirroredStrategy):
            tolerance = 0.0001
        keras.backend.set_image_data_format('channels_last')
        np.random.seed(_RANDOM_SEED)
        random_seed.set_random_seed(_RANDOM_SEED)
        num_samples = 9984
        x_train = np.random.rand(num_samples, 1)
        y_train = (3 * x_train)
        x_train = x_train.astype('float32')
        y_train = y_train.astype('float32')
        x_predict = [[1.0], [2.0], [3.0], [4.0]]
        model = keras.Sequential()
        model.add(keras.layers.Dense(10, activation='relu', input_shape=(1,)))
        model.add(keras.layers.Dense(10, activation='relu'))
        model.add(keras.layers.Dense(10, activation='relu'))
        model.add(keras.layers.Dense(1))
        initial_weights = model.get_weights()

        def fit_and_predict(with_distribution=None):
            model.set_weights(initial_weights)
            model.compile(loss=keras.losses.mean_squared_error, optimizer=gradient_descent.GradientDescentOptimizer(0.5), distribute=with_distribution)
            (training_inputs, eval_inputs, predict_inputs) = get_correctness_test_inputs(use_numpy, with_distribution, x_train, y_train, x_predict)
            model.fit(**training_inputs)
            eval_result = model.evaluate(**eval_inputs)
            weights = model.get_weights()
            predict_result = model.predict(**predict_inputs)
            return (weights, eval_result, predict_result)
        (wts_with_ds, eval_with_ds, predict_with_ds) = fit_and_predict(with_distribution=distribution)
        (wts_without_ds, eval_without_ds, predict_without_ds) = fit_and_predict(with_distribution=None)
        self.assertAllClose(wts_with_ds, wts_without_ds, atol=tolerance, rtol=tolerance)
        self.assertAllClose(eval_with_ds, eval_without_ds, atol=tolerance, rtol=tolerance)
        self.assertAllClose(predict_with_ds, predict_without_ds, atol=tolerance, rtol=tolerance)