model must be created under scope of diststrat it will be trained with forward pass compute gradients update weights runs a single training step on a batch of data return self trainstepdata if not self runeagerly kwargs jitcompile self jitcompile if self supportsreduceretracing kwargs updatereduceretracing true onestepondata tf functiononestepondata kwargs tf autograph experimental donotconvert def onesteponiteratoriterator runs a single test step on a batch of data return self teststepdata if not self runeagerly and self jitcompile kwargs jitcompile true if self supportsreduceretracing kwargs updatereduceretracing true onestepondata tf functiononestepondata kwargs tf autograph experimental donotconvert def onesteponiteratoriterator runs a predict test step on a batch of data return self predictstepdata if not self runeagerly and self jitcompile kwargs jitcompile true if self supportsreduceretracing kwargs updatereduceretracing true onestepondata tf functiononestepondata kwargs tf autograph experimental donotconvert def onestepondatadistributeddata data data0 outputs self distributestrategy run onestepondata argsdata outputs reduceperreplica outputs self distributestrategy reductionself distributereductionmethod return outputs tf autograph experimental donotconvert def multistepondatadata outputs onestepondatadistributeddata 1 for singlestepdata in data1 stepoutputs onestepondatadistributedsinglestepdata outputs tf nest mapstructure lambda t1 t2 concatt1 t2 outputs stepoutputs return outputs if self stepsperexecution 1 predictfunction multistepondata else predictfunction onestepondatadistributed if not self runeagerly kwargs if self supportsreduceretracing kwargs updatereduceretracing true predictfunction tf functionpredictfunction kwargs self predictfunction predictfunction tracebackutils filtertraceback def fit self xnone ynone batchsizenone epochs1 verboseauto callbacksnone validationsplit0 0 validationdatanone shuffletrue classweightnone sampleweightnone initialepoch0 stepsperepochnone validationstepsnone validationbatchsizenone validationfreq1 self assertcompilecalledfit todo respect compiled trainable state self evalepochiterator none if validationsplit and validationdata is none create the validation data using the training data only supported for tfnumpyjax arrays x y sampleweight validationdata dataadapterutils trainvalidationsplit x y sampleweight validationsplitvalidationsplit if validationdata valx valy valsampleweight dataadapterutils unpackxysampleweightvalidationdata create an iterator that yields batches for one epoch epochiterator tfepochiterator xx yy sampleweightsampleweight batchsizebatchsize stepsperepochstepsperepoch shuffleshuffle classweightclassweight distributestrategyself distributestrategy stepsperexecutionself stepsperexecution container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochsepochs stepsepochiterator numbatches modelself self stoptraining false self maketrainfunction callbacks ontrainbegin traininglogs none logs none for epoch in rangeinitialepoch epochs self resetmetrics callbacks onepochbeginepoch with epochiterator catchstopiteration for step iterator in epochiterator enumerateepoch callbacks ontrainbatchbeginstep logs self trainfunctioniterator callbacks ontrainbatchend step self pythonifylogslogs if self stoptraining break override with model metrics instead of last step logs epochlogs self getmetricsresult run validation if validationdata and self shouldevalepoch validationfreq create epochiterator for evaluation and cache it if getattrself evalepochiterator none is none self evalepochiterator tfepochiterator xvalx yvaly sampleweightvalsampleweight batchsizevalidationbatchsize or batchsize distributestrategyself distributestrategy stepsperexecutionself stepsperexecution stepsperepochvalidationsteps shufflefalse vallogs self evaluate xvalx yvaly sampleweightvalsampleweight batchsizevalidationbatchsize or batchsize stepsvalidationsteps callbackscallbacks returndicttrue usecachedevaldatasettrue vallogs val name val for name val in vallogs items epochlogs updatevallogs callbacks onepochendepoch epochlogs traininglogs epochlogs if self stoptraining break if isinstanceself optimizer optimizersmodule optimizer and epochs 0 self optimizer finalizevariablevaluesself trainableweights if evalepochiterator exists delete it after all epochs are done if getattrself evalepochiterator none is not none del self evalepochiterator callbacks ontrainendlogstraininglogs return self history tracebackutils filtertraceback def evaluate self xnone ynone batchsizenone verboseauto sampleweightnone stepsnone callbacksnone returndictfalse kwargs self assertcompilecalledevaluate todo respect compiled trainable state usecachedevaldataset kwargs popusecachedevaldataset false if kwargs raise valueerrorfarguments not recognized kwargs if usecachedevaldataset epochiterator self evalepochiterator else create an iterator that yields batches of inputtarget data epochiterator tfepochiterator xx yy sampleweightsampleweight batchsizebatchsize stepsperepochsteps shufflefalse distributestrategyself distributestrategy stepsperexecutionself stepsperexecution container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochs1 stepsepochiterator numbatches modelself self maketestfunction self stopevaluating false callbacks ontestbegin logs none self resetmetrics with epochiterator catchstopiteration for step iterator in epochiterator enumerateepoch callbacks ontestbatchbeginstep logs self testfunctioniterator callbacks ontestbatchendstep self pythonifylogslogs if self stopevaluating break logs self getmetricsresult callbacks ontestendlogs if returndict return logs return self flattenmetricsinorderlogs tracebackutils filtertraceback def predict self x batchsizenone verboseauto stepsnone callbacksnone create an iterator that yields batches of input data epochiterator tfepochiterator xx batchsizebatchsize stepsperepochsteps shufflefalse distributestrategyself distributestrategy stepsperexecutionself stepsperexecution container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochs1 stepsepochiterator numbatches modelself def appendtooutputsbatchoutputs outputs if outputs is none outputs tf nest mapstructure lambda batchoutput batchoutput batchoutputs else tree mapstructureupto batchoutputs lambda output batchoutput output appendbatchoutput outputs batchoutputs return outputs def getdataiterator suppress the error when still have remaining data reraise the error for tfepochiterator catchstopiteration to catch when no data left backwards compatibility shims catches errors when an iterator runs out of data try yield self tfsync except stopiteration tf errors outofrangeerror if self numbatches is none self numbatches self stepsseen warnings warn your input ran out of data interrupting training make sure that your dataset or generator can generate at least stepsperepoch epochs batches you may need to use the repeat function when building your dataset stacklevel2 self currentiterator none self dataadapter onepochend def reduceperreplicavalues strategy reduction if reduction auto reduction sum ignore tpu strategy which should default to first def reducev order perreplica objects for collectiveallreducestrategy and concat replicas strategy gatherv axis0 v might not have the same shape on different replicas if isperreplicainstancev shapes tf concat tf expanddimstf shapesinglevalue0 axis0 for singlevalue in v values axis0 allshapes strategy gathershapes axis0 else v is a tensor this may happen when say we have 2x1 multiworker allshapes strategy gather tf expanddimstf shapev0 axis0 axis0 replicas tf split replicas numorsizesplitsallshapes numstrategy numreplicasinsync orderedreplicas numreplicasperworker lenstrategy extended workerdevices for replicaid in rangenumreplicasperworker orderedreplicas replicasreplicaid numreplicasperworker return concatorderedreplicas def concattensors axis0 correctly order tpu perreplica objects replicas strategy experimentallocalresultsv when distributed datasets are created from tensors numpy tpustrategy experimentaldistributedataset shards data in replica host order and tpustrategy experimentallocalresults returns it in host replica order numreplicasperhost strategy extended numreplicasperhost orderedreplicas for replicaid in rangenumreplicasperhost orderedreplicas replicasreplicaid numreplicasperhost return concatorderedreplicas def collectiveallreducemultiworkerstrategy return isinstancestrategy tf distribute multiworkermirroredstrategy and strategy extended inmultiworkermode def isperreplicainstanceobj return isinstanceobj tf distribute distributedvalues and isinstance obj tf internal compositetensor def isscalarx return isinstancex tf tensor tf variable and x shape rank 0 def istpumultihoststrategy return istpustrategystrategy and strategy extended numhosts 1 def istpustrategystrategy return istpustrategyclassstrategy class def istpustrategyclassclz def istpustratk return k name startswithtpustrategy if istpustratclz return true return anymapistpustrategyclass clz bases def converttonpifnotraggedx if isinstancex tf raggedtensor return x return x numpy def potentiallyraggedconcattensors if lentensors 1 return tensors0 elif isinstancetensors0 tf sparsetensor return tf sparse concataxis0 spinputstensors elif isinstancetensors0 tf raggedtensor return tf concattensors axis0 nonbatchshapes tf stacktf shapetensor1 for tensor in tensors constantdims tf math reduceall nonbatchshapes nonbatchshapes 1 axis0 if tf math reduceallconstantdims numpy item all nonbatch dims are constant if isscalartensors0 return tf stacktensors axis0 else return tf concattensors axis0 first identify constant inner dimensions by finding the rightmost dimension that is not constant constantinnerdimensions constantdims numpy tolist 1 indexfalse if there are constant inner dimensions define a constant inner shape if constantinnerdimensions 0 constantinnershape none else constantinnershape tensors0 shapeconstantinnerdimensions return tf ragged constant tensor numpy for tensor in tensors innershapeconstantinnershape mergedims0 1 model must be created under scope of diststrat it will be trained with forward pass compute gradients update weights runs a single training step on a batch of data runs a single training step given a dataset iterator runs a single test step on a batch of data runs a single test step given a dataset iterator runs a predict test step on a batch of data todo respect compiled trainable state create the validation data using the training data only supported for tf numpy jax arrays create an iterator that yields batches for one epoch container that configures and calls callbacks override with model metrics instead of last step logs run validation create epochiterator for evaluation and cache it if _eval_epoch_iterator exists delete it after all epochs are done todo respect compiled trainable state create an iterator that yields batches of input target data container that configures and calls callbacks create an iterator that yields batches of input data container that configures and calls callbacks returns data for the next execution suppress the error when still have remaining data re raise the error for tfepochiterator catch_stop_iteration to catch when no data left backwards compatibility shims catches errors when an iterator runs out of data attempt to reduce the structure values to single values given values a tf tensor or a perreplica structure which represents the values across all the replicas reduce_per_replica attempts to reduce those values and returns the corresponding structure that represents only single values currently reduce_per_replica is only used for reducing the metric results from tf distribute strategy run depending on the underlying strategy implementation values may be a perreplica object which can be thought of as a collection of values across the replicas or a tf tensor if the strategy has already conducted the reduction for the downstream library there are five possible outcomes of reduction 1 if the values is a structure of simple tf tensor s meaning that reduction is not actually needed reduce_per_replica returns the structure as is 2 else if reduction auto then the best reduction strategy is chosen based on the current environment this should only be used for training cases fit 3 else if reduction first then reduce_per_replica returns the values of the first replica this is used in the case of training and evaluation where values is expected to hold the same value across the replicas as a result of strategy s synchronization across the replicas reduce_per_replica does not synchronize the values 4 else if reduction sum then reduce_per_replica returns the sum of values for all replicas this may be used in the custom training loop case where each replica contain different values which are not synchronized 5 else if reduction concat then reduce_per_replica returns the concatenation of the values across the replicas along the axis of dimension 0 this is used in the inference case predict args values structure of perreplica objects or tf tensor s tf tensor s are returned as is strategy tf distribute strategy object reduction one of auto first concat or sum auto will select first when used under a tpustrategy or sum otherwise returns structure of tensor s representing the result of reduction raises valueerror if the reduction method is not supported ignore tpu strategy which should default to first reduce a single perreplica object order perreplica objects for collectiveallreducestrategy and concat v might not have the same shape on different replicas v is a tensor this may happen when say we have 2x1 multi worker concats tensor s along axis correctly order tpu perreplica objects when distributed datasets are created from tensors numpy tpustrategy experimental_distribute_dataset shards data in replica host order and tpustrategy experimental_local_results returns it in host replica order concats tensor s along their first dimension args tensors list of tensor s returns concatenation of the inputs along the first dimension of type np ndarray if all input shapes are compatible or tf raggedtensor if not all non batch dims are constant first identify constant inner dimensions by finding the rightmost dimension that is not constant if there are constant inner dimensions define a constant inner shape
import contextlib
import warnings

import numpy as np
import tensorflow as tf
import tree
from packaging.version import Version
from tensorflow.python.eager import context as tf_context

from keras import callbacks as callbacks_module
from keras import metrics as metrics_module
from keras import optimizers as optimizers_module
from keras.trainers import trainer as base_trainer
from keras.trainers.data_adapters import data_adapter_utils
from keras.trainers.epoch_iterator import EpochIterator
from keras.utils import traceback_utils


class TensorFlowTrainer(base_trainer.Trainer):
    def __init__(self):
        super().__init__()
        self.train_function = None
        self.test_function = None
        self.predict_function = None

        
        
        if tf.distribute.has_strategy():
            self._distribute_strategy = tf.distribute.get_strategy()
        else:
            self._distribute_strategy = None

        self._distribute_reduction_method = None
        self._supports_reduce_retracing = Version(tf.__version__) >= Version(
            "2.9.0"
        )

    @property
    def distribute_strategy(self):
        return self._distribute_strategy or tf.distribute.get_strategy()

    @property
    def distribute_reduction_method(self):
        return self._distribute_reduction_method or "auto"

    @distribute_reduction_method.setter
    def distribute_reduction_method(self, value):
        self._distribute_reduction_method = value

    def train_step(self, data):
        x, y, sample_weight = data_adapter_utils.unpack_x_y_sample_weight(data)

        
        with tf.GradientTape() as tape:
            if self._call_has_training_arg:
                y_pred = self(x, training=True)
            else:
                y_pred = self(x)
            loss = self.compute_loss(
                x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
            )
            self._loss_tracker.update_state(loss)
            if self.optimizer is not None:
                loss = self.optimizer.scale_loss(loss)

        
        if self.trainable_weights:
            trainable_weights = self.trainable_weights
            gradients = tape.gradient(loss, trainable_weights)

            
            self.optimizer.apply_gradients(zip(gradients, trainable_weights))
        else:
            warnings.warn("The model does not have any trainable weights.")

        return self.compute_metrics(x, y, y_pred, sample_weight=sample_weight)

    def test_step(self, data):
        x, y, sample_weight = data_adapter_utils.unpack_x_y_sample_weight(data)
        if self._call_has_training_arg:
            y_pred = self(x, training=False)
        else:
            y_pred = self(x)
        loss = self.compute_loss(
            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
        )
        self._loss_tracker.update_state(loss)
        return self.compute_metrics(x, y, y_pred, sample_weight=sample_weight)

    def predict_step(self, data):
        x, _, _ = data_adapter_utils.unpack_x_y_sample_weight(data)
        if self._call_has_training_arg:
            y_pred = self(x, training=False)
        else:
            y_pred = self(x)
        return y_pred

    def make_train_function(self, force=False):
        if self.train_function is not None and not force:
            return self.train_function

        @tf.autograph.experimental.do_not_convert
        def one_step_on_data(data):
            
            return self.train_step(data)

        if not self.run_eagerly:
            kwargs = {"jit_compile": self.jit_compile}
            if self._supports_reduce_retracing:
                kwargs.update({"reduce_retracing": True})
            one_step_on_data = tf.function(one_step_on_data, **kwargs)

        @tf.autograph.experimental.do_not_convert
        def one_step_on_iterator(iterator):
            
            data = next(iterator)
            outputs = self.distribute_strategy.run(
                one_step_on_data, args=(data,)
            )
            outputs = reduce_per_replica(
                outputs,
                self.distribute_strategy,
                reduction=self.distribute_reduction_method,
            )
            return outputs

        @tf.autograph.experimental.do_not_convert
        def multi_step_on_iterator(iterator):
            for _ in range(self.steps_per_execution):
                outputs = one_step_on_iterator(iterator)
            return outputs

        if self.steps_per_execution > 1:
            train_function = multi_step_on_iterator
        else:
            train_function = one_step_on_iterator

        if not self.run_eagerly:
            kwargs = {}
            if self._supports_reduce_retracing:
                kwargs.update({"reduce_retracing": True})
            train_function = tf.function(train_function, **kwargs)

        self.train_function = train_function

    def make_test_function(self, force=False):
        if self.test_function is not None and not force:
            return self.test_function

        @tf.autograph.experimental.do_not_convert
        def one_step_on_data(data):
            
            return self.test_step(data)

        if not self.run_eagerly and self.jit_compile:
            kwargs = {"jit_compile": True}
            if self._supports_reduce_retracing:
                kwargs.update({"reduce_retracing": True})
            one_step_on_data = tf.function(one_step_on_data, **kwargs)

        @tf.autograph.experimental.do_not_convert
        def one_step_on_iterator(iterator):
            
            data = next(iterator)
            outputs = self.distribute_strategy.run(
                one_step_on_data, args=(data,)
            )
            outputs = reduce_per_replica(
                outputs,
                self.distribute_strategy,
                reduction=self.distribute_reduction_method,
            )
            return outputs

        @tf.autograph.experimental.do_not_convert
        def multi_step_on_iterator(iterator):
            for _ in range(self.steps_per_execution):
                outputs = one_step_on_iterator(iterator)
            return outputs

        if self.steps_per_execution > 1:
            test_function = multi_step_on_iterator
        else:
            test_function = one_step_on_iterator

        if not self.run_eagerly:
            kwargs = {}
            if self._supports_reduce_retracing:
                kwargs.update({"reduce_retracing": True})
            test_function = tf.function(test_function, **kwargs)

        self.test_function = test_function

    def make_predict_function(self, force=False):
        if self.predict_function is not None and not force:
            return self.predict_function

        @tf.autograph.experimental.do_not_convert
        def one_step_on_data(data):
            
            return self.predict_step(data)

        if not self.run_eagerly and self.jit_compile:
            kwargs = {"jit_compile": True}
            if self._supports_reduce_retracing:
                kwargs.update({"reduce_retracing": True})
            one_step_on_data = tf.function(one_step_on_data, **kwargs)

        @tf.autograph.experimental.do_not_convert
        def one_step_on_data_distributed(data):
            data = data[0]
            outputs = self.distribute_strategy.run(
                one_step_on_data, args=(data,)
            )
            outputs = reduce_per_replica(
                outputs,
                self.distribute_strategy,
                reduction=self.distribute_reduction_method,
            )
            return outputs

        @tf.autograph.experimental.do_not_convert
        def multi_step_on_data(data):
            outputs = one_step_on_data_distributed(data[:1])
            for single_step_data in data[1:]:
                step_outputs = one_step_on_data_distributed([single_step_data])
                outputs = tf.nest.map_structure(
                    lambda t1, t2: concat([t1, t2]), outputs, step_outputs
                )
            return outputs

        if self.steps_per_execution > 1:
            predict_function = multi_step_on_data
        else:
            predict_function = one_step_on_data_distributed

        if not self.run_eagerly:
            kwargs = {}
            if self._supports_reduce_retracing:
                kwargs.update({"reduce_retracing": True})

            predict_function = tf.function(predict_function, **kwargs)

        self.predict_function = predict_function

    @traceback_utils.filter_traceback
    def fit(
        self,
        x=None,
        y=None,
        batch_size=None,
        epochs=1,
        verbose="auto",
        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_batch_size=None,
        validation_freq=1,
    ):
        self._assert_compile_called("fit")
        
        self._eval_epoch_iterator = None
        if validation_split and validation_data is None:
            
            
            (
                x,
                y,
                sample_weight,
            ), validation_data = data_adapter_utils.train_validation_split(
                (x, y, sample_weight), validation_split=validation_split
            )

        if validation_data:
            (
                val_x,
                val_y,
                val_sample_weight,
            ) = data_adapter_utils.unpack_x_y_sample_weight(validation_data)

        
        epoch_iterator = TFEpochIterator(
            x=x,
            y=y,
            sample_weight=sample_weight,
            batch_size=batch_size,
            steps_per_epoch=steps_per_epoch,
            shuffle=shuffle,
            class_weight=class_weight,
            distribute_strategy=self.distribute_strategy,
            steps_per_execution=self.steps_per_execution,
        )

        
        if not isinstance(callbacks, callbacks_module.CallbackList):
            callbacks = callbacks_module.CallbackList(
                callbacks,
                add_history=True,
                add_progbar=verbose != 0,
                verbose=verbose,
                epochs=epochs,
                steps=epoch_iterator.num_batches,
                model=self,
            )

        self.stop_training = False
        self.make_train_function()
        callbacks.on_train_begin()
        training_logs = None
        logs = None
        for epoch in range(initial_epoch, epochs):
            self.reset_metrics()
            callbacks.on_epoch_begin(epoch)
            with epoch_iterator.catch_stop_iteration():
                for step, iterator in epoch_iterator.enumerate_epoch():
                    callbacks.on_train_batch_begin(step)
                    logs = self.train_function(iterator)
                    callbacks.on_train_batch_end(
                        step, self._pythonify_logs(logs)
                    )
                    if self.stop_training:
                        break

            
            epoch_logs = self.get_metrics_result()

            
            if validation_data and self._should_eval(epoch, validation_freq):
                
                if getattr(self, "_eval_epoch_iterator", None) is None:
                    self._eval_epoch_iterator = TFEpochIterator(
                        x=val_x,
                        y=val_y,
                        sample_weight=val_sample_weight,
                        batch_size=validation_batch_size or batch_size,
                        distribute_strategy=self.distribute_strategy,
                        steps_per_execution=self.steps_per_execution,
                        steps_per_epoch=validation_steps,
                        shuffle=False,
                    )
                val_logs = self.evaluate(
                    x=val_x,
                    y=val_y,
                    sample_weight=val_sample_weight,
                    batch_size=validation_batch_size or batch_size,
                    steps=validation_steps,
                    callbacks=callbacks,
                    return_dict=True,
                    _use_cached_eval_dataset=True,
                )
                val_logs = {
                    "val_" + name: val for name, val in val_logs.items()
                }
                epoch_logs.update(val_logs)

            callbacks.on_epoch_end(epoch, epoch_logs)
            training_logs = epoch_logs
            if self.stop_training:
                break

        if (
            isinstance(self.optimizer, optimizers_module.Optimizer)
            and epochs > 0
        ):
            self.optimizer.finalize_variable_values(self.trainable_weights)

        
        if getattr(self, "_eval_epoch_iterator", None) is not None:
            del self._eval_epoch_iterator
        callbacks.on_train_end(logs=training_logs)
        return self.history

    @traceback_utils.filter_traceback
    def evaluate(
        self,
        x=None,
        y=None,
        batch_size=None,
        verbose="auto",
        sample_weight=None,
        steps=None,
        callbacks=None,
        return_dict=False,
        **kwargs,
    ):
        self._assert_compile_called("evaluate")
        
        use_cached_eval_dataset = kwargs.pop("_use_cached_eval_dataset", False)
        if kwargs:
            raise ValueError(f"Arguments not recognized: {kwargs}")

        if use_cached_eval_dataset:
            epoch_iterator = self._eval_epoch_iterator
        else:
            
            epoch_iterator = TFEpochIterator(
                x=x,
                y=y,
                sample_weight=sample_weight,
                batch_size=batch_size,
                steps_per_epoch=steps,
                shuffle=False,
                distribute_strategy=self.distribute_strategy,
                steps_per_execution=self.steps_per_execution,
            )

        
        if not isinstance(callbacks, callbacks_module.CallbackList):
            callbacks = callbacks_module.CallbackList(
                callbacks,
                add_history=True,
                add_progbar=verbose != 0,
                verbose=verbose,
                epochs=1,
                steps=epoch_iterator.num_batches,
                model=self,
            )

        self.make_test_function()
        self.stop_evaluating = False
        callbacks.on_test_begin()
        logs = None
        self.reset_metrics()
        with epoch_iterator.catch_stop_iteration():
            for step, iterator in epoch_iterator.enumerate_epoch():
                callbacks.on_test_batch_begin(step)
                logs = self.test_function(iterator)
                callbacks.on_test_batch_end(step, self._pythonify_logs(logs))
                if self.stop_evaluating:
                    break
        logs = self.get_metrics_result()
        callbacks.on_test_end(logs)

        if return_dict:
            return logs
        return self._flatten_metrics_in_order(logs)

    @traceback_utils.filter_traceback
    def predict(
        self, x, batch_size=None, verbose="auto", steps=None, callbacks=None
    ):
        
        epoch_iterator = TFEpochIterator(
            x=x,
            batch_size=batch_size,
            steps_per_epoch=steps,
            shuffle=False,
            distribute_strategy=self.distribute_strategy,
            steps_per_execution=self.steps_per_execution,
        )

        
        if not isinstance(callbacks, callbacks_module.CallbackList):
            callbacks = callbacks_module.CallbackList(
                callbacks,
                add_history=True,
                add_progbar=verbose != 0,
                verbose=verbose,
                epochs=1,
                steps=epoch_iterator.num_batches,
                model=self,
            )

        def append_to_outputs(batch_outputs, outputs):
            if outputs is None:
                outputs = tf.nest.map_structure(
                    lambda batch_output: [batch_output],
                    batch_outputs,
                )
            else:
                tree.map_structure_up_to(
                    batch_outputs,
                    lambda output, batch_output: output.append(batch_output),
                    outputs,
                    batch_outputs,
                )
            return outputs

        def get_data(iterator):
            
            data = []
            for _ in range(self.steps_per_execution):
                try:
                    single_step_data = next(iterator)
                except (StopIteration, tf.errors.OutOfRangeError) as e:
                    if hasattr(data, "__len__") and len(data) > 0:
                        
                        return data
                    else:
                        
                        
                        
                        raise e
                data.append(single_step_data)
            return data

        self.make_predict_function()
        self.stop_predicting = False
        callbacks.on_predict_begin()
        outputs = None
        with epoch_iterator.catch_stop_iteration():
            for step, iterator in epoch_iterator.enumerate_epoch():
                callbacks.on_predict_batch_begin(step)
                data = get_data(iterator)
                batch_outputs = self.predict_function(data)
                outputs = append_to_outputs(batch_outputs, outputs)
                callbacks.on_predict_batch_end(step, {"outputs": batch_outputs})
                if self.stop_predicting:
                    break
        callbacks.on_predict_end()
        outputs = tree.map_structure_up_to(
            batch_outputs, potentially_ragged_concat, outputs
        )
        return tf.nest.map_structure(convert_to_np_if_not_ragged, outputs)

    def train_on_batch(
        self,
        x,
        y=None,
        sample_weight=None,
        class_weight=None,
        return_dict=False,
    ):
        self._assert_compile_called("train_on_batch")
        self.make_train_function()
        if class_weight is not None:
            if sample_weight is not None:
                raise ValueError(
                    "Arguments `sample_weight` and `class_weight` "
                    "cannot be specified at the same time. "
                    f"Received: sample_weight={sample_weight}, "
                    f"class_weight={class_weight}"
                )
            sample_weight = data_adapter_utils.class_weight_to_sample_weights(
                y, class_weight
            )

        def data():
            yield (x, y, sample_weight)

        logs = self.train_function(data())
        logs = tf.nest.map_structure(lambda x: np.array(x), logs)
        if return_dict:
            return logs
        return self._flatten_metrics_in_order(logs)

    def test_on_batch(
        self,
        x,
        y=None,
        sample_weight=None,
        return_dict=False,
    ):
        self._assert_compile_called("test_on_batch")
        self.make_test_function()

        def data():
            yield (x, y, sample_weight)

        logs = self.test_function(data())
        logs = tf.nest.map_structure(lambda x: np.array(x), logs)
        if return_dict:
            return logs
        return self._flatten_metrics_in_order(logs)

    def predict_on_batch(self, x):
        self.make_predict_function()
        batch_outputs = self.predict_function([(x,)])
        batch_outputs = tf.nest.map_structure(
            convert_to_np_if_not_ragged, batch_outputs
        )
        return batch_outputs

    
    @property
    def compiled_metrics(self):
        class DeprecatedCompiledMetric:
            def update_state(_, y, y_pred, sample_weight=None):
                return self._compiled_metrics_update_state(
                    y, y_pred, sample_weight=sample_weight
                )

        return DeprecatedCompiledMetric()

    def _compiled_metrics_update_state(self, y, y_pred, sample_weight=None):
        warnings.warn(
            "`model.compiled_metrics()` is deprecated. "
            "Instead, use e.g.:\n"
            "```\n"
            "for metric in self.metrics:\n"
            "    metric.update_state(y, y_pred)\n"
            "```\n",
            stacklevel=2,
        )
        for metric in self.metrics:
            if isinstance(metric, metrics_module.Mean):
                metric.update_state(y_pred, sample_weight=sample_weight)
            else:
                metric.update_state(y, y_pred, sample_weight=sample_weight)

    def compiled_loss(
        self, y, y_pred, sample_weight=None, regularization_losses=None
    ):
        warnings.warn(
            "`model.compiled_loss()` is deprecated. "
            "Instead, use `model.compute_loss(x, y, y_pred, sample_weight)`.",
        )
        return self.compute_loss(
            x=None, y=y, y_pred=y_pred, sample_weight=sample_weight
        )

    def loss(self, y, y_pred, sample_weight=None):
        warnings.warn(
            "`model.loss` is deprecated. "
            "Instead, use `model.compute_loss(x, y, y_pred, sample_weight)`.",
        )
        return self.compute_loss(
            x=None, y=y, y_pred=y_pred, sample_weight=sample_weight
        )


class TFEpochIterator(EpochIterator):
    def __init__(self, distribute_strategy=None, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._distribute_strategy = distribute_strategy
        dataset = self.data_adapter.get_tf_dataset()
        if not isinstance(dataset, tf.distribute.DistributedDataset):
            dataset = self._distribute_strategy.experimental_distribute_dataset(
                dataset
            )
        self._distributed_dataset = dataset
        self._steps_seen = 0

    def enumerate_epoch(self):
        if self.steps_per_epoch:
            if not self._current_iterator:
                self._current_iterator = iter(self._distributed_dataset)
            for step in range(
                0, self.steps_per_epoch, self.steps_per_execution
            ):
                yield step, self._current_iterator
        else:
            iterator = iter(self._distributed_dataset)
            if self.num_batches:
                for step in range(
                    0, self.num_batches, self.steps_per_execution
                ):
                    yield step, iterator
            else:
                step = -1
                while True:
                    step += self.steps_per_execution
                    self._steps_seen = step + 1
                    yield step, iterator
        self.data_adapter.on_epoch_end()

    def tf_sync(self):
        tf_context.async_wait()

    @contextlib.contextmanager
    def catch_stop_iteration(self):
        
        try:
            yield
            self.tf_sync()
        except (StopIteration, tf.errors.OutOfRangeError):
            if self._num_batches is None:
                self._num_batches = self._steps_seen
            warnings.warn(
                "Your input ran out of data; interrupting training. "
                "Make sure that your dataset or generator can generate "
                "at least `steps_per_epoch * epochs` batches. "
                "You may need to use the `.repeat()` "
                "function when building your dataset.",
                stacklevel=2,
            )
            self._current_iterator = None
            self.data_adapter.on_epoch_end()


def reduce_per_replica(values, strategy, reduction):
    

    if reduction == "auto":
        reduction = "sum"  

    def _reduce(v):
        
        if _collective_all_reduce_multi_worker(strategy):
            if reduction == "concat":
                return _multi_worker_concat(v, strategy)
            elif reduction == "sum":
                return strategy.reduce("SUM", v, axis=None)

        if not _is_per_replica_instance(v):
            return v
        elif reduction == "first":
            return strategy.experimental_local_results(v)[0]
        elif reduction == "concat":
            if _is_tpu_multi_host(strategy):
                return _tpu_multi_host_concat(v, strategy)
            else:
                return concat(strategy.experimental_local_results(v))
        elif reduction == "sum":
            return tf.reduce_sum(strategy.experimental_local_results(v))
        else:
            raise ValueError(
                '`reduction` must be "first", "concat", "sum", or "auto". '
                f"Received: reduction={reduction}."
            )

    return tf.nest.map_structure(_reduce, values)


def _multi_worker_concat(v, strategy):
    
    replicas = strategy.gather(v, axis=0)
    
    if _is_per_replica_instance(v):
        shapes = tf.concat(
            [
                tf.expand_dims(tf.shape(single_value)[0], axis=0)
                for single_value in v.values
            ],
            axis=0,
        )
        all_shapes = strategy.gather(shapes, axis=0)
    else:
        
        all_shapes = strategy.gather(
            tf.expand_dims(tf.shape(v)[0], axis=0), axis=0
        )

    replicas = tf.split(
        replicas,
        num_or_size_splits=all_shapes,
        num=strategy.num_replicas_in_sync,
    )
    ordered_replicas = []
    num_replicas_per_worker = len(strategy.extended.worker_devices)
    for replica_id in range(num_replicas_per_worker):
        ordered_replicas += replicas[replica_id::num_replicas_per_worker]
    return concat(ordered_replicas)


def concat(tensors, axis=0):
    
    if isinstance(tensors[0], tf.SparseTensor):
        return tf.sparse.concat(axis=axis, sp_inputs=tensors)
    elif _is_scalar(tensors[0]):
        return tf.stack(tensors, axis=axis)
    else:
        return tf.concat(tensors, axis=axis)


def _tpu_multi_host_concat(v, strategy):
    
    replicas = strategy.experimental_local_results(v)
    
    
    
    
    num_replicas_per_host = strategy.extended.num_replicas_per_host
    ordered_replicas = []
    for replica_id in range(num_replicas_per_host):
        ordered_replicas += replicas[replica_id::num_replicas_per_host]
    return concat(ordered_replicas)


def _collective_all_reduce_multi_worker(strategy):
    return (
        isinstance(strategy, tf.distribute.MultiWorkerMirroredStrategy)
    ) and strategy.extended._in_multi_worker_mode()


def _is_per_replica_instance(obj):
    return isinstance(obj, tf.distribute.DistributedValues) and isinstance(
        obj, tf.__internal__.CompositeTensor
    )


def _is_scalar(x):
    return isinstance(x, (tf.Tensor, tf.Variable)) and x.shape.rank == 0


def _is_tpu_multi_host(strategy):
    return _is_tpu_strategy(strategy) and strategy.extended.num_hosts > 1


def _is_tpu_strategy(strategy):
    return _is_tpu_strategy_class(strategy.__class__)


def _is_tpu_strategy_class(clz):
    def is_tpu_strat(k):
        return k.__name__.startswith("TPUStrategy")

    if is_tpu_strat(clz):
        return True
    return any(map(_is_tpu_strategy_class, clz.__bases__))


def convert_to_np_if_not_ragged(x):
    if isinstance(x, tf.RaggedTensor):
        return x
    return x.numpy()


def potentially_ragged_concat(tensors):
    
    if len(tensors) == 1:
        return tensors[0]
    elif isinstance(tensors[0], tf.SparseTensor):
        return tf.sparse.concat(axis=0, sp_inputs=tensors)
    elif isinstance(tensors[0], tf.RaggedTensor):
        return tf.concat(tensors, axis=0)

    non_batch_shapes = tf.stack([tf.shape(tensor)[1:] for tensor in tensors])
    constant_dims = tf.math.reduce_all(
        non_batch_shapes == non_batch_shapes[:1], axis=0
    )
    if tf.math.reduce_all(constant_dims).numpy().item():
        
        if _is_scalar(tensors[0]):
            return tf.stack(tensors, axis=0)
        else:
            return tf.concat(tensors, axis=0)

    
    
    constant_inner_dimensions = (
        constant_dims.numpy().tolist()[::-1].index(False)
    )
    
    if constant_inner_dimensions == 0:
        constant_inner_shape = None
    else:
        constant_inner_shape = tensors[0].shape[-constant_inner_dimensions:]
    return tf.ragged.constant(
        [tensor.numpy() for tensor in tensors], inner_shape=constant_inner_shape
    ).merge_dims(0, 1)