# Copyright 2018 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.
# ==============================================================================
"""Part of the Keras training engine related to distributed training."""

import numpy as np
import tensorflow.compat.v2 as tf

from keras import backend
from keras import callbacks as cbks
from keras.distribute import distribute_coordinator_utils as dc
from keras.distribute import distributed_training_utils_v1 as dist_utils
from keras.engine import partial_batch_padding_handler as padding_util
from keras.engine import training_arrays_v1
from keras.engine import training_utils_v1
from keras.utils.generic_utils import Progbar
from keras.utils.mode_keys import ModeKeys

# isort: off
from tensorflow.python.distribute import input_lib
from tensorflow.python.platform import tf_logging as logging


def _per_replica_execution_function(model, mode):
    exec_func = model._make_execution_function(mode)
    return (
        exec_func.inputs,
        exec_func.outputs,
        exec_func.updates_op,
        exec_func.session_kwargs,
    )


def _build_model(strategy, model, mode, inputs, targets=None):
    if model._compile_distribution:
        dist_utils.clone_model_on_replicas(
            model, strategy, mode, inputs=inputs, targets=targets
        )
    else:
        dist_utils._build_distributed_network(
            model, strategy, mode, inputs, targets
        )


def _make_train_step_fn(model, mode, strategy, output_labels):
    """Create step fn.

    Args:
      model: a Keras Model instance.
      mode: One of ModeKeys.TRAIN/ModeKeys.TEST/ModeKeys.PREDICT.
      strategy: a `tf.distribute.Strategy` instance.
      output_labels: the output labels for the step function.

    Returns:
      A step function to run by `tf.distribute.Strategy`.
    """

    def _step_fn(ctx, inputs):
        """A step fn that returns update ops."""
        if isinstance(inputs, (tuple, list)) and len(inputs) == 2:
            inputs, targets = inputs
        else:
            targets = None

        # When input feature is a dictionary of tensors, dictionary is
        # flattended to an array and passed as a model input. This results in
        # input mismatch when model input layer names are not sorted in
        # alphabetical order as `nest.flatten()`sorts dictionary elements by
        # keys. As so, transform input tensors into an array and order it along
        # `model._feed_input_names`.
        if isinstance(inputs, dict):
            inputs = [
                inputs[input_name] for input_name in model._feed_input_names
            ]

        _build_model(strategy, model, mode, inputs, targets)

        (
            grouped_inputs,
            grouped_outputs,
            grouped_updates,
            grouped_session_args,
        ) = strategy.extended.call_for_each_replica(
            _per_replica_execution_function,
            args=(dist_utils.get_distributed_model(model, mode), mode),
        )
        (
            all_inputs,
            all_outputs,
            all_updates,
            all_session_args,
        ) = dist_utils.unwrap_values(
            strategy,
            grouped_inputs,
            grouped_outputs,
            grouped_updates,
            grouped_session_args,
        )
        combined_fn = backend.function(
            all_inputs,
            all_outputs,
            updates=all_updates,
            name="distributed_" + str(mode) + "_function",
            **all_session_args
        )

        for label, output in zip(output_labels, combined_fn.outputs):
            if label == "loss":
                reduce_op = tf.distribute.ReduceOp.SUM
            else:
                # We reduce all other metrics using mean for now. This is
                # temporary workaround until new metrics are in place.
                reduce_op = tf.distribute.ReduceOp.MEAN
            ctx.set_last_step_output(label, output, reduce_op)

        # TODO(priyag, sourabhbajaj): Ignoring these things from the
        # combined_fn: feed_dict, session kwargs, run options, run_metadata for
        # now. These should be handled appropriately
        return combined_fn.updates_op

    return _step_fn


def experimental_tpu_fit_loop(
    model,
    dataset,
    epochs=100,
    verbose=1,
    callbacks=None,
    initial_epoch=0,
    steps_per_epoch=None,
    val_dataset=None,
    validation_steps=None,
    validation_freq=1,
):
    """Fit loop for training with TPU tf.distribute.Strategy.

    Args:
        model: Keras Model instance.
        dataset: Dataset that returns inputs and targets
        epochs: Number of times to iterate over the data
        verbose: Integer, Verbosity mode, 0, 1 or 2
        callbacks: List of callbacks to be called during training
        initial_epoch: Epoch at which to start training
            (useful for resuming a previous training run)
        steps_per_epoch: Total number of steps (batches of samples)
            before declaring one epoch finished and starting the
            next epoch. Ignored with the default value of `None`.
        val_dataset: Dataset for validation data.
        validation_steps: Number of steps to run validation for
            (only if doing validation from data tensors).
            Ignored with the default value of `None`.
        validation_freq: Only relevant if validation data is provided. Integer
            or `collections.abc.Container` instance (e.g. list, tuple, etc.). If
            an integer, specifies how many training epochs to run before a new
            validation run is performed, e.g. `validation_freq=2` runs
            validation every 2 epochs. If a Container, specifies the epochs on
            which to run validation, e.g. `validation_freq=[1, 2, 10]` runs
            validation at the end of the 1st, 2nd, and 10th epochs.

    Returns:
        Returns `None`.

    Raises:
        ValueError: in case of invalid arguments.
    """
    mode = ModeKeys.TRAIN

    current_strategy = model._distribution_strategy
    iteration_value = min(
        steps_per_epoch, current_strategy.extended.steps_per_run
    )
    steps_per_run = backend.variable(
        value=iteration_value, dtype="int32", name="steps_per_run"
    )

    # TODO(fchollet): add support for `steps_per_epoch=None` in TPU loops.
    iterator = dist_utils.get_iterator(dataset, current_strategy)

    scope = dist_utils.distributed_scope(
        strategy=current_strategy, learning_phase=1
    )
    scope.__enter__()

    out_labels = model.metrics_names or []

    step_fn = _make_train_step_fn(
        model, ModeKeys.TRAIN, current_strategy, out_labels
    )

    # Add initial dummy values for loss and other metric tensors.
    initial_loop_values = {}
    initial_loop_values["loss"] = tf.constant(1e7)
    for m in model._get_training_eval_metrics():
        tensor = m.result()
        initial_loop_values[m.name] = tf.zeros(tensor.shape, tensor.dtype)

    ctx = current_strategy.extended.experimental_run_steps_on_iterator(
        step_fn,
        iterator,
        iterations=steps_per_run,
        initial_loop_values=initial_loop_values,
    )
    train_op = ctx.run_op
    output_tensors = ctx.last_step_outputs

    do_validation = bool(validation_steps)

    if model._compile_distribution:
        dist_utils._copy_weights_to_distributed_model(model, mode)

    callbacks = cbks.configure_callbacks(
        callbacks,
        model,
        do_validation=do_validation,
        epochs=epochs,
        steps_per_epoch=steps_per_epoch,
        verbose=verbose,
        count_mode="steps",
        mode=mode,
    )

    # Calculate the steps each time on the device.
    steps_to_run = [current_strategy.extended.steps_per_run] * (
        steps_per_epoch // current_strategy.extended.steps_per_run
    )
    if steps_per_epoch % current_strategy.extended.steps_per_run:
        steps_to_run.append(
            steps_per_epoch % current_strategy.extended.steps_per_run
        )
    target_steps = len(steps_to_run)

    callbacks._call_begin_hook(mode)

    initial_epoch = model._maybe_load_initial_epoch_from_ckpt(
        initial_epoch, mode
    )

    for epoch in range(initial_epoch, epochs):
        dist_utils._reset_metrics(model)
        callbacks.on_epoch_begin(epoch)
        epoch_logs = {}
        step_index = 0
        prev_step_count = None
        current_step = 0
        while current_step < target_steps:
            step_count = steps_to_run[current_step]
            batch_logs = {
                "batch": step_index,
                "size": 1,
                "num_steps": step_count,
            }
            callbacks._call_batch_hook(mode, "begin", step_index, batch_logs)
            if prev_step_count is None or step_count != prev_step_count:
                backend.get_session().run(steps_per_run.assign(step_count))
                prev_step_count = step_count
            try:
                _, outputs = backend.batch_get_value([train_op, output_tensors])
            except tf.errors.OutOfRangeError:
                logging.warning(
                    "Your dataset iterator ran out of data; "
                    "interrupting training. Make sure that your dataset "
                    "can generate at least `steps_per_epoch * epochs` "
                    "batches (in this case, %d batches)."
                    % steps_per_epoch
                    * epochs
                )
                break

            batch_logs.update(outputs)
            callbacks._call_batch_hook(mode, "end", step_index, batch_logs)
            step_index = step_index + step_count
            current_step += 1

            if callbacks.model.stop_training:
                break

        if do_validation and training_utils_v1.should_run_validation(
            validation_freq, epoch
        ):
            logging.info("Running validation at fit epoch: %s", epoch)

            if model._compile_distribution:
                # Since we create a new clone from the original model we need to
                # copy the weights back to the original model before we can run
                # validation.
                dist_utils._copy_weights_to_original_model(
                    model, ModeKeys.TRAIN
                )

            val_outs = experimental_tpu_test_loop(
                model,
                val_dataset,
                steps=validation_steps,
                verbose=verbose,
                callbacks=callbacks,
            )
            if not isinstance(val_outs, list):
                val_outs = [val_outs]
            # Same labels assumed.
            for label, val_out in zip(out_labels, val_outs):
                epoch_logs["val_" + label] = val_out

        callbacks.on_epoch_end(epoch, epoch_logs)
        if callbacks.model.stop_training:
            break
    model._successful_loop_finish = True
    callbacks._call_end_hook(mode)

    if model._compile_distribution:
        # Copy the weights back from the replicated model to the original model.
        dist_utils._copy_weights_to_original_model(model, ModeKeys.TRAIN)
    scope.__exit__(None, None, None)
    return model.history


def experimental_tpu_test_loop(
    model, dataset, verbose=0, steps=None, callbacks=None
):
    """Test loop for evaluating with TPU tf.distribute.Strategy.

    Args:
        model: Keras Model instance.
        dataset: Dataset for input data.
        verbose: Integer, Verbosity mode 0 or 1.
        steps: Total number of steps (batches of samples)
            before declaring predictions finished.
            Ignored with the default value of `None`.
        callbacks: List of callbacks to be called during training

    Returns:
        Scalar loss (if the model has a single output and no metrics)
        or list of scalars (if the model has multiple outputs
        and/or metrics). The attribute `model.metrics_names` will give you
        the display labels for the outputs.
    """
    mode = ModeKeys.TEST
    current_strategy = model._distribution_strategy
    iterator = dist_utils.get_iterator(dataset, current_strategy)

    scope = dist_utils.distributed_scope(
        strategy=current_strategy, learning_phase=0
    )
    scope.__enter__()

    out_labels = model.metrics_names

    def _test_step_fn(inputs):
        """A fn that returns output of single test step."""
        if isinstance(inputs, (tuple, list)) and len(inputs) == 2:
            inputs, targets = inputs
        else:
            targets = None

        (
            tf.distribute.get_replica_context().merge_call(
                _build_model, args=(model, mode, inputs, targets)
            )
        )

        (_, outputs, updates, _) = _per_replica_execution_function(
            dist_utils.get_distributed_model(model, mode), mode
        )
        with tf.control_dependencies([updates]):
            return [tf.identity(out) for out in outputs]

    test_input_data = iterator.get_next()
    per_replica_outputs = current_strategy.run(
        _test_step_fn, args=(test_input_data,)
    )
    output_tensors = {}
    for label, output in zip(out_labels, per_replica_outputs):
        if label == "loss":
            reduce_op = tf.distribute.ReduceOp.SUM
        else:
            # We reduce all other metrics using mean for now. This is temporary
            # workaround until new metrics are in place.
            reduce_op = tf.distribute.ReduceOp.MEAN
        output_tensors[label] = current_strategy.reduce(
            reduce_op, output, axis=None
        )
    test_op = tf.group(list(output_tensors.values()))

    if verbose >= 1:
        progbar = Progbar(target=steps)

    if model._compile_distribution:
        dist_utils._copy_weights_to_distributed_model(model, mode)

    dist_utils._reset_metrics(model)

    callbacks = cbks.configure_callbacks(
        callbacks,
        model,
        do_validation=False,
        epochs=1,
        steps_per_epoch=steps,
        verbose=verbose,
        count_mode="steps",
        mode=ModeKeys.TEST,
    )
    callbacks._call_begin_hook(mode)

    outs = [0.0] * len(model.metrics_names)
    if steps is not None:
        target_steps = steps
    else:
        raise ValueError(
            "Number of steps could not be inferred from the data, "
            "please pass the steps argument."
        )

    current_step = 0
    while current_step < target_steps:
        batch_logs = {"batch": current_step, "size": 1}
        callbacks._call_batch_hook(mode, "begin", current_step, batch_logs)
        try:
            _, batch_outs = backend.batch_get_value([test_op, output_tensors])
        except tf.errors.OutOfRangeError:
            warning_msg = (
                "Make sure that your dataset can generate at least "
                "`steps` batches (in this case, {} batches).".format(steps)
            )

            logging.warning(
                "Your dataset iterator ran out of data; "
                "interrupting evaluation. " + warning_msg
            )
            target_steps = current_step
            break
        for i, label in enumerate(model.metrics_names):
            if i == 0:
                # Loss is stateless metrics.
                outs[i] += batch_outs[label]
            else:
                # For all stateful metrics, the aggregation is handled by
                # mirrored vars.
                outs[i] = batch_outs[label]

        batch_logs = callbacks.make_logs(model, batch_logs, outs, mode)
        callbacks._call_batch_hook(mode, "end", current_step, batch_logs)
        if verbose == 1:
            progbar.update(current_step + 1)
        current_step += 1

    if verbose >= 1:
        # Progress bar finishes at the end.
        progbar.update(target_steps)
    callbacks._call_end_hook(mode)

    scope.__exit__(None, None, None)
    if len(outs) > 0:
        outs[0] /= target_steps

    if len(outs) == 1:
        return outs[0]
    return outs


def experimental_tpu_predict_loop(
    model, dataset, verbose=0, steps=None, callbacks=None
):
    """Predict loop for predicting with TPU tf.distribute.Strategy.

    Args:
        model: Keras Model instance.
        dataset: Dataset for input data.
        verbose: Integer, Verbosity mode 0 or 1.
        steps: Total number of steps (batches of samples)
            before declaring `_predict_loop` finished.
            Ignored with the default value of `None`.
        callbacks: List of callbacks to be called during training

    Returns:
        Array of predictions (if the model has a single output)
        or list of arrays of predictions
        (if the model has multiple outputs).
    """
    mode = ModeKeys.PREDICT
    dataset_fully_shaped = dist_utils.is_dataset_shape_fully_defined(dataset)
    padding_handler = None
    if not dataset_fully_shaped:
        # TODO(hongjunchoi): Investigate whether operations from
        # PartialBatchPaddingHandler are unnecessarily pruned out
        # during graph optimization.
        padding_handler = padding_util.PartialBatchPaddingHandler(
            model._feed_output_shapes
        )
        batch_size, _, prefetch_buffer = input_lib._get_dataset_attributes(
            dataset
        )
        padding_handler.padded_batch_size = batch_size
        padding_handler.padding_mask = dataset.reduce(
            padding_handler.padding_mask, padding_handler.update_mask
        )

        dataset = dataset.map(padding_handler.pad_batch)
        dataset = dataset.unbatch()
        # Upon this point, it is guaranteed that the dataset does not
        # have partial batches. Thus, we set `drop_remainder=True` to
        # get static shape information about the elements in the dataset.
        dataset = dataset.batch(batch_size, drop_remainder=True)

        if prefetch_buffer is not None:
            dataset = dataset.prefetch(prefetch_buffer)

    current_strategy = model._distribution_strategy
    iterator = dist_utils.get_iterator(dataset, current_strategy)

    scope = dist_utils.distributed_scope(
        strategy=current_strategy, learning_phase=0
    )
    scope.__enter__()

    def _predict_step_fn(inputs):
        """A fn that returns output of single prediction step."""

        (
            tf.distribute.get_replica_context().merge_call(
                _build_model, args=(model, mode, inputs)
            )
        )

        (_, outputs, updates, _) = _per_replica_execution_function(
            dist_utils.get_distributed_model(model, mode), mode
        )

        with tf.control_dependencies([updates]):
            return [tf.identity(out) for out in outputs]

    # TODO(hongjunchoi): When numpy array is passed as an input to `predict()`
    # use numpy arrays directly to avoid cumulating unnecessary input pipeline
    # ops.
    predict_input_data = iterator.get_next()
    per_replica_outputs = current_strategy.run(
        _predict_step_fn, args=(predict_input_data,)
    )
    output_tensors = dist_utils.flatten_per_replica_values(
        current_strategy, per_replica_outputs
    )

    if verbose >= 1:
        progbar = Progbar(target=steps)

    if model._compile_distribution:
        dist_utils._copy_weights_to_distributed_model(model, mode)

    dist_utils._reset_metrics(model)

    callbacks = cbks.configure_callbacks(
        callbacks,
        model,
        do_validation=False,
        epochs=1,
        steps_per_epoch=steps,
        verbose=verbose,
        count_mode="steps",
        mode=mode,
    )
    callbacks._call_begin_hook(mode)

    # Since we do not know how many samples we will see, we cannot pre-allocate
    # the returned Numpy arrays. Instead, we store one array per batch seen
    # and concatenate them upon returning.
    num_model_outputs = len(model.output_names)
    unconcatenated_outs = [[] for _ in range(num_model_outputs)]
    if steps is not None:
        target_steps = steps
    else:
        raise ValueError(
            "Number of steps could not be inferred from the data, "
            "please pass the steps argument."
        )

    current_step = 0
    while current_step < target_steps:
        batch_logs = {"batch": current_step, "size": 1}
        callbacks._call_batch_hook(mode, "begin", current_step, batch_logs)
        try:
            predict_ops = tf.group(output_tensors)
            _, batch_outs = backend.batch_get_value(
                [predict_ops, output_tensors]
            )

        except tf.errors.OutOfRangeError:
            warning_msg = (
                "Make sure that your dataset can generate at least "
                "`steps` batches (in this case, {} batches).".format(steps)
            )

            logging.warning(
                "Your dataset iterator ran out of data; "
                "interrupting evaluation. " + warning_msg
            )
            break

        # TODO(priyag): maybe need to unwrap the outputs first for
        # MirroredStrategy.
        for i in range(num_model_outputs):
            output_start_index = i * current_strategy.num_replicas_in_sync
            output_end_index = (
                output_start_index + current_strategy.num_replicas_in_sync
            )
            single_model_output = batch_outs[
                output_start_index:output_end_index
            ]
            unconcatenated_outs[i].extend(single_model_output)

        batch_logs = callbacks.make_logs(model, batch_logs, batch_outs, mode)
        callbacks._call_batch_hook(mode, "end", current_step, batch_logs)
        if verbose == 1:
            progbar.update(current_step + 1)
        current_step += 1

    if verbose >= 1:
        # Progress bar finishes at the end.
        progbar.update(current_step)

    callbacks._call_end_hook(mode)

    scope.__exit__(None, None, None)

    if len(unconcatenated_outs) == 1:
        prediction_result = np.concatenate(unconcatenated_outs[0], axis=0)
    else:
        prediction_result = [
            np.concatenate(out, axis=0) for out in unconcatenated_outs
        ]

    if padding_handler:
        prediction_result = padding_handler.apply_mask(prediction_result)

    return prediction_result


class DistributionSingleWorkerTrainingLoop(training_utils_v1.TrainingLoop):
    """Training loop for distribution strategy with single worker."""

    def fit(
        self,
        model,
        x=None,
        y=None,
        batch_size=None,
        epochs=1,
        verbose=1,
        callbacks=None,
        validation_split=0.0,
        validation_data=None,
        shuffle=True,
        class_weight=None,
        sample_weight=None,
        initial_epoch=0,
        steps_per_epoch=None,
        validation_steps=None,
        validation_freq=1,
        **kwargs
    ):
        """Fit loop for Distribution Strategies."""
        dist_utils.validate_callbacks(
            input_callbacks=callbacks, optimizer=model.optimizer
        )
        dist_utils.validate_inputs(x, y)

        batch_size, steps_per_epoch = dist_utils.process_batch_and_step_size(
            model._distribution_strategy,
            x,
            batch_size,
            steps_per_epoch,
            ModeKeys.TRAIN,
            validation_split=validation_split,
        )
        batch_size = model._validate_or_infer_batch_size(
            batch_size, steps_per_epoch, x
        )
        dataset = model._distribution_standardize_user_data(
            x,
            y,
            sample_weight=sample_weight,
            class_weight=class_weight,
            batch_size=batch_size,
            validation_split=validation_split,
            shuffle=shuffle,
            epochs=epochs,
        )
        if not dist_utils.is_distributing_by_cloning(model):
            with model._distribution_strategy.scope():
                (dataset, _, _) = model._standardize_user_data(
                    dataset,
                    sample_weight=sample_weight,
                    class_weight=class_weight,
                    batch_size=batch_size,
                    validation_split=validation_split,
                    shuffle=shuffle,
                )

        val_dataset = None
        if validation_data:
            (
                val_x,
                val_y,
                val_sample_weights,
            ) = training_utils_v1.unpack_validation_data(validation_data)
            dist_utils.validate_inputs(val_x, val_y)
            _, validation_steps = dist_utils.process_batch_and_step_size(
                model._distribution_strategy,
                val_x,
                batch_size,
                validation_steps,
                ModeKeys.TEST,
            )

            val_dataset = model._distribution_standardize_user_data(
                val_x,
                val_y,
                sample_weight=val_sample_weights,
                class_weight=None,
                batch_size=batch_size,
                validation_split=validation_split,
                shuffle=shuffle,
                allow_partial_batch=True,
            )
        elif validation_split:
            raise ValueError(
                "validation_split argument is not supported with "
                "distribution strategies."
            )

        if backend.is_tpu_strategy(model._distribution_strategy):
            steps_per_epoch = training_utils_v1.infer_steps_for_dataset(
                model,
                dataset,
                steps_per_epoch,
                epochs,
                steps_name="steps_per_epoch",
            )
            if steps_per_epoch is None:
                raise ValueError(
                    "Number of steps could not be inferred from the data, "
                    "please pass the steps_per_epoch argument."
                )

            if not tf.executing_eagerly():
                # Run TPU training in a custom loop in graph mode.
                return experimental_tpu_fit_loop(
                    model,
                    dataset,
                    epochs=epochs,
                    verbose=verbose,
                    callbacks=callbacks,
                    val_dataset=val_dataset,
                    initial_epoch=initial_epoch,
                    steps_per_epoch=steps_per_epoch,
                    validation_steps=validation_steps,
                    validation_freq=validation_freq,
                )

        return training_arrays_v1.fit_loop(
            model,
            dataset,
            batch_size=batch_size,
            epochs=epochs,
            verbose=verbose,
            callbacks=callbacks,
            val_inputs=val_dataset,
            shuffle=shuffle,
            initial_epoch=initial_epoch,
            steps_per_epoch=steps_per_epoch,
            validation_steps=validation_steps,
            validation_freq=validation_freq,
            steps_name="steps_per_epoch",
        )

    def evaluate(
        self,
        model,
        x=None,
        y=None,
        batch_size=None,
        verbose=1,
        sample_weight=None,
        steps=None,
        callbacks=None,
        **kwargs
    ):
        """Evaluate loop for Distribution Strategies."""
        dist_utils.validate_inputs(x, y)
        batch_size, steps = dist_utils.process_batch_and_step_size(
            model._distribution_strategy, x, batch_size, steps, ModeKeys.TEST
        )
        batch_size = model._validate_or_infer_batch_size(batch_size, steps, x)
        dataset = model._distribution_standardize_user_data(
            x,
            y,
            sample_weight=sample_weight,
            batch_size=batch_size,
            allow_partial_batch=True,
        )

        if backend.is_tpu_strategy(model._distribution_strategy):
            steps = training_utils_v1.infer_steps_for_dataset(
                model, dataset, steps, steps_name="steps"
            )
            if steps is None:
                raise ValueError(
                    "Number of steps could not be inferred from the data, "
                    "please pass the steps argument."
                )

            if not tf.executing_eagerly():
                # Run TPU evaluation in a custom loop in graph mode.
                return experimental_tpu_test_loop(
                    model,
                    dataset,
                    verbose=verbose,
                    steps=steps,
                    callbacks=callbacks,
                )

        return training_arrays_v1.test_loop(
            model,
            inputs=dataset,
            batch_size=batch_size,
            verbose=verbose,
            steps=steps,
            callbacks=callbacks,
        )

    def predict(
        self,
        model,
        x,
        batch_size=None,
        verbose=0,
        steps=None,
        callbacks=None,
        **kwargs
    ):
        """Predict loop for Distribution Strategies."""
        dist_utils.validate_inputs(x=x, y=None)
        batch_size, steps = dist_utils.process_batch_and_step_size(
            model._distribution_strategy, x, batch_size, steps, ModeKeys.PREDICT
        )
        batch_size = model._validate_or_infer_batch_size(batch_size, steps, x)
        dataset = model._distribution_standardize_user_data(
            x, batch_size=batch_size, allow_partial_batch=True
        )
        if backend.is_tpu_strategy(model._distribution_strategy):
            steps = training_utils_v1.infer_steps_for_dataset(
                model, dataset, steps, steps_name="steps"
            )
            if steps is None:
                raise ValueError(
                    "Number of steps could not be inferred from the data, "
                    "please pass the steps argument."
                )
            if not tf.executing_eagerly():
                return experimental_tpu_predict_loop(
                    model,
                    dataset,
                    verbose=verbose,
                    steps=steps,
                    callbacks=callbacks,
                )
        return training_arrays_v1.predict_loop(
            model,
            dataset,
            batch_size=batch_size,
            verbose=verbose,
            steps=steps,
            callbacks=callbacks,
        )


def _train_with_multi_worker(method):
    """Decorator handles multi worker training with distribution strategy."""

    def wrapper(model, **kwargs):
        def _worker_fn(_):
            callbacks = kwargs.pop("callbacks", None)
            filtered_callbacks = dist_utils.filter_distributed_callbacks(
                callbacks, model
            )
            kwargs["callbacks"] = filtered_callbacks
            return method(model, **kwargs)

        return dc.run_distribute_coordinator(
            _worker_fn, model._distribution_strategy
        )

    return wrapper


class DistributionMultiWorkerTrainingLoop(training_utils_v1.TrainingLoop):
    """Training loop for distribution strategy with multiple worker."""

    def __init__(self, single_worker_loop):
        self._single_worker_loop = single_worker_loop

    def fit(self, *args, **kwargs):
        return _train_with_multi_worker(self._single_worker_loop.fit)(
            *args, **kwargs
        )

    def evaluate(self, *args, **kwargs):
        return _train_with_multi_worker(self._single_worker_loop.evaluate)(
            *args, **kwargs
        )

    def predict(self, *args, **kwargs):
        # Currently predict is still using the single worker implementation.
        return self._single_worker_loop.predict(*args, **kwargs)
