this method is stateless and is intended for use with jax grad kwargs if self callhastrainingarg kwargstraining training ypred nontrainablevariables losses self statelesscall trainablevariables nontrainablevariables x returnlossestrue kwargs trainablemapping zipself trainablevariables trainablevariables with backend statelessscopestatemappingtrainablemapping note that this is needed for the regularization loss which need the latest value of trainnontrainable variables loss self computeloss x y ypred sampleweight allowemptytrue if losses loss ops sumlosses unscaledloss loss if training and self optimizer is not none scale loss with a statelessscope to use an update scale variable mapping listzipself optimizer variables optimizervariables with backend statelessscopestatemappingmapping loss self optimizer scalelossloss return loss unscaledloss ypred nontrainablevariables def eagerbuildself databatch modelunbuilt not alllayer built for layer in self flattenlayers compilemetricsunbuilt self compilemetrics is not none and not self compilemetrics built if modelunbuilt or compilemetricsunbuilt def convertdatatospecd if d is none return none return backend kerastensord shape d dtype dataspec tree mapstructureconvertdatatospec databatch xspec yspec sampleweightspec dataadapterutils unpackxysampleweightdataspec note that this call run the forward path and trigger variable creation ypredspec backend computeoutputspecself call xspec if compilemetricsunbuilt this will trigger the metric variable creation backend computeoutputspec self computemetrics xspec yspec ypredspec sampleweightsampleweightspec if self optimizer is not none and not self optimizer built build optimizer self optimizer buildself trainablevariables def trainstepself state data trainablevariables nontrainablevariables optimizervariables metricsvariables state x y sampleweight dataadapterutils unpackxysampleweightdata gradfn jax valueandgrad self computelossandupdates hasauxtrue loss aux grads gradfn trainablevariables nontrainablevariables x y sampleweight trainingtrue optimizervariablesoptimizervariables unscaledloss ypred nontrainablevariables aux trainablevariables optimizervariables self optimizer statelessapply optimizervariables grads trainablevariables with backend statelessscope statemapping refv v for refv v in zipself metricsvariables metricsvariables as scope self losstracker updatestateunscaledloss logs self computemetricsx y ypred sampleweight newmetricsvariables for refv in self metricsvariables newv scope getcurrentvaluerefv if newv is none newv refv value newmetricsvariables appendnewv metricsvariables newmetricsvariables state self enforcejaxstatesharding trainablevariables nontrainablevariables optimizervariables metricsvariables return logs state def teststepself state data trainablevariables nontrainablevariables metricsvariables state x y sampleweight dataadapterutils unpackxysampleweightdata loss aux self computelossandupdates trainablevariables nontrainablevariables x y sampleweight trainingfalse unscaledloss ypred nontrainablevariables aux with backend statelessscope statemapping refv v for refv v in zipself metricsvariables metricsvariables as scope self losstracker updatestateunscaledloss logs self computemetricsx y ypred sampleweight newmetricsvariables for refv in self metricsvariables newv scope getcurrentvaluerefv if newv is none newv refv value newmetricsvariables appendnewv metricsvariables newmetricsvariables trainablevariables nontrainablevariables metricsvariables self enforcejaxstatesharding trainablevariablestrainablevariables nontrainablevariablesnontrainablevariables optimizervariablesnone metricsvariablesmetricsvariables state trainablevariables nontrainablevariables metricsvariables return logs state def predictstepself state data trainablevariables nontrainablevariables state kwargs if self callhastrainingarg kwargstraining false x dataadapterutils unpackxysampleweightdata outputs nontrainablevariables self statelesscall trainablevariables nontrainablevariables x kwargs trainablevariables nontrainablevariables self enforcejaxstatesharding trainablevariablestrainablevariables nontrainablevariablesnontrainablevariables optimizervariablesnone metricsvariablesnone return outputs trainablevariables nontrainablevariables def maketrainfunctionself forcefalse if self trainfunction is not none and not force return self trainfunction def onetrainstepstate data data data0 return self trainstepstate data def multitrainstepsstate data for singlestepdata in data logs state onetrainstepstate singlestepdata return logs state if self stepsperexecution 1 trainstep multitrainsteps else trainstep onetrainstep if not self runeagerly and self jitcompile note that we mark the state and data to be donated to jax so that jax will reuse the memory buffer for outputs this will reduce the memory usage of the training function by half partialjax jit donateargnamesstate def compiledtrainstepstate data return trainstepstate data self trainfunction compiledtrainstep else self trainfunction trainstep def maketestfunctionself forcefalse if self testfunction is not none and not force return self testfunction def oneteststepstate data data data0 return self teststepstate data def multiteststepsstate data for singlestepdata in data logs state oneteststepstate singlestepdata return logs state if self stepsperexecution 1 teststep multiteststeps else teststep oneteststep if not self runeagerly and self jitcompile note that we mark the state and data to be donated to jax so that jax will reuse the memory buffer for outputs this will reduce the memory usage of the training function by half partialjax jit donateargnamesstate def compiledteststepstate data return teststepstate data self testfunction compiledteststep else self testfunction teststep def makepredictfunctionself forcefalse if self predictfunction is not none and not force return self predictfunction def onepredictstepstate data data data0 return self predictstepstate data def multipredictstepsstate data outputs state onepredictstepstate data 1 for singlestepdata in data1 stepoutputs state onepredictstep state singlestepdata outputs tree mapstructure lambda t1 t2 jax numpy concatenatet1 t2 outputs stepoutputs return outputs state if self stepsperexecution 1 predictstep multipredictsteps else predictstep onepredictstep if not self runeagerly and self jitcompile jax jit def compiledpredictstepstate data return predictstepstate data self predictfunction compiledpredictstep else self predictfunction predictstep 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 epochiterator xx yy sampleweightsampleweight batchsizebatchsize stepsperepochstepsperepoch shuffleshuffle classweightclassweight stepsperexecutionself stepsperexecution needsbuilding not alllayer built for layer in self flattenlayers or not self optimizer built or self compilemetrics is not none and not self compilemetrics built if needsbuilding build the model on one batch of data for data in epochiterator enumerateepochreturntypenp databatch data0 self eagerbuilddatabatch break container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochsepochs stepsepochiterator numbatches modelself self recordtrainingstateshardingspec self maketrainfunction self stoptraining false callbacks ontrainbegin for epoch in rangeinitialepoch epochs self resetmetrics callbacks onepochbeginepoch trainablevariables v value for v in self trainablevariables nontrainablevariables v value for v in self nontrainablevariables optimizervariables v value for v in self optimizer variables metricsvariables v value for v in self metricsvariables self purgemodelvariables for step data in epochiterator enumerateepochreturntypenp callbacks callbacks ontrainbatchbeginstep train step state trainablevariables nontrainablevariables optimizervariables metricsvariables data self distributedatadata logs state self trainfunctionstate data trainablevariables nontrainablevariables optimizervariables metricsvariables state setting jaxstate enables callbacks to force a state sync if they need to self jaxstate trainablevariables trainablevariables nontrainablevariables nontrainablevariables optimizervariables optimizervariables metricsvariables metricsvariables callbacks callbacks ontrainbatchendstep self pythonifylogslogs if self stoptraining break reattach state to model variables note doing this after each step would be a big performance bottleneck self jaxstatesync override with model metrics instead of last step logs the jax spmdmode is need for multiprocess context since the metrics values are replicated and we don t want to do a all gather and only need the local copy of the value with jax spmdmodeallowall 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 epochiterator xvalx yvaly sampleweightvalsampleweight batchsizevalidationbatchsize or batchsize 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 self jaxstate none 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 epochiterator xx yy sampleweightsampleweight batchsizebatchsize stepsperepochsteps shufflefalse stepsperexecutionself stepsperexecution needsbuilding not all layer built for layer in self flattenlayers or self compilemetrics is not none and not self compilemetrics built if needsbuilding build the model on one batch of data for data in epochiterator enumerateepochreturntypenp databatch data0 self eagerbuilddatabatch break container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochs1 stepsepochiterator numbatches modelself self recordtrainingstateshardingspec self maketestfunction self stopevaluating false callbacks ontestbegin logs none self resetmetrics trainablevariables v value for v in self trainablevariables nontrainablevariables v value for v in self nontrainablevariables metricsvariables v value for v in self metricsvariables self purgemodelvariablesoptimizervariablesfalse for step data in epochiterator enumerateepochreturntypenp callbacks ontestbatchbeginstep state trainablevariables nontrainablevariables metricsvariables data self distributedatadata logs state self testfunctionstate data trainablevariables nontrainablevariables metricsvariables state setting jaxstate enables callbacks to force a state sync if they need to self jaxstate i wouldn t recommend modifying nontrainable model state during evaluate but it s allowed trainablevariables trainablevariables nontrainablevariables nontrainablevariables metricsvariables metricsvariables callbacks ontestbatchendstep self pythonifylogslogs if self stopevaluating break reattach state back to model self jaxstatesync the jax spmdmode is need for multiprocess context since the metrics values are replicated and we don t want to do a all gather and only need the local copy of the value with jax spmdmodeallowall logs self getmetricsresult callbacks ontestendlogs self jaxstate none 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 epochiterator xx batchsizebatchsize stepsperepochsteps shufflefalse stepsperexecutionself stepsperexecution if not alllayer built for layer in self flattenlayers build the model on one batch of data for data in epochiterator enumerateepochreturntypenp build model x dataadapterutils unpackxysampleweightdata0 with backend statelessscope selfx break container that configures and calls callbacks if not isinstancecallbacks callbacksmodule callbacklist callbacks callbacksmodule callbacklist callbacks addhistorytrue addprogbarverbose 0 verboseverbose epochs1 stepsepochiterator numbatches modelself self recordtrainingstateshardingspec self makepredictfunction self stoppredicting false callbacks onpredictbegin def appendtooutputsbatchoutputs outputs if outputs is none outputs tree mapstructure lambda batchoutput batchoutput batchoutputs else tree mapstructureupto batchoutputs lambda output batchoutput output appendbatchoutput outputs batchoutputs return outputs trainablevariables v value for v in self trainablevariables nontrainablevariables v value for v in self nontrainablevariables state trainablevariables nontrainablevariables outputs none for step x in epochiterator enumerateepochreturntypenp callbacks onpredictbatchbeginstep x self distributedatax batchoutputs state self predictfunctionstate x outputs appendtooutputsbatchoutputs outputs callbacks onpredictbatchendstep outputs batchoutputs if self stoppredicting break callbacks onpredictend return tree mapstructureuptobatchoutputs np concatenate outputs def trainonbatch self x ynone sampleweightnone classweightnone returndictfalse self assertcompilecalledtrainonbatch if classweight is not none if sampleweight is not none raise valueerror arguments sampleweight and classweight cannot be specified at the same time freceived sampleweightsampleweight fclassweightclassweight sampleweight dataadapterutils classweighttosampleweights y classweight data x y sampleweight data self distributedatadata maybe build model self eagerbuilddata self recordtrainingstateshardingspec self maketrainfunction train step trainablevariables v value for v in self trainablevariables nontrainablevariables v value for v in self nontrainablevariables optimizervariables v value for v in self optimizer variables metricsvariables v value for v in self metricsvariables state trainablevariables nontrainablevariables optimizervariables metricsvariables logs state self trainfunctionstate data state sync trainablevariables nontrainablevariables optimizervariables metricsvariables state self jaxstate trainablevariables trainablevariables nontrainablevariables nontrainablevariables optimizervariables optimizervariables metricsvariables metricsvariables self jaxstatesync format return values logs tree mapstructurelambda x np arrayx logs if returndict return logs return self flattenmetricsinorderlogs def testonbatch self x ynone sampleweightnone returndictfalse self assertcompilecalledtestonbatch data x y sampleweight data self distributedatadata maybe build model self eagerbuilddata self recordtrainingstateshardingspec self maketestfunction test step trainablevariables v value for v in self trainablevariables nontrainablevariables v value for v in self nontrainablevariables metricsvariables v value for v in self metricsvariables state trainablevariables nontrainablevariables metricsvariables logs state self testfunctionstate data state sync trainablevariables nontrainablevariables metricsvariables state self jaxstate trainablevariables trainablevariables nontrainablevariables nontrainablevariables metricsvariables metricsvariables self jaxstatesync format return values logs tree mapstructurelambda x np arrayx logs if returndict return logs return self flattenmetricsinorderlogs def predictonbatchself x if not alllayer built for layer in self flattenlayers build model with backend statelessscope selfx self recordtrainingstateshardingspec self makepredictfunction trainablevariables v value for v in self trainablevariables nontrainablevariables v value for v in self nontrainablevariables state trainablevariables nontrainablevariables batchoutputs state self predictfunctionstate x batchoutputs tree mapstructurelambda x np arrayx batchoutputs return batchoutputs def jaxstatesyncself if not getattrself jaxstate none return trainablevariables self jaxstate gettrainablevariables none nontrainablevariables self jaxstate get nontrainablevariables none optimizervariables self jaxstate getoptimizervariables none metricsvariables self jaxstate getmetricsvariables none if trainablevariables for refv v in zipself trainablevariables trainablevariables refv assignv if nontrainablevariables for refv v in zip self nontrainablevariables nontrainablevariables refv assignv if optimizervariables for refv v in zipself optimizer variables optimizervariables refv assignv if metricsvariables for refv v in zipself metricsvariables metricsvariables refv assignv def distributedataself data distribution distributionlib distribution if distribution is not none def distributesinglevalued layout distribution getdatalayoutd shape return jaxdistributionlib distributedatainputd layout return jax treeutil treemapdistributesinglevalue data else return data def recordtrainingstateshardingspecself self trainablevariableshardings v value sharding for v in self trainablevariables self nontrainablevariableshardings v value sharding for v in self nontrainablevariables if hasattrself optimizer and self optimizer is not none self optimizervariableshardings v value sharding for v in self optimizer variables else self optimizervariableshardings self metricsvariableshardings v value sharding for v in self metricsvariables def enforcejaxstatesharding self trainablevariablesnone nontrainablevariablesnone optimizervariablesnone metricsvariablesnone trainablevariables trainablevariables or nontrainablevariables nontrainablevariables or optimizervariables optimizervariables or metricsvariables metricsvariables or for i in rangelentrainablevariables trainablevariablesi jax lax withshardingconstraint trainablevariablesi self trainablevariableshardingsi for i in rangelennontrainablevariables nontrainablevariablesi jax lax withshardingconstraint nontrainablevariablesi self nontrainablevariableshardingsi for i in rangelenoptimizervariables optimizervariablesi jax lax withshardingconstraint optimizervariablesi self optimizervariableshardingsi for i in rangelenmetricsvariables metricsvariablesi jax lax withshardingconstraint metricsvariablesi self metricsvariableshardingsi return trainablevariables nontrainablevariables optimizervariables metricsvariables def purgemodelvariables self trainablevariablestrue nontrainablevariablestrue optimizervariablestrue metricvariablestrue if trainablevariables for v in self trainablevariables v value none if nontrainablevariables for v in self nontrainablevariables v value none if optimizervariables for v in self optimizer variables v value none if metricvariables for v in self metricsvariables v value none this method is stateless and is intended for use with jax grad note that this is needed for the regularization loss which need the latest value of train non trainable variables scale loss with a statelessscope to use an update scale variable note that this __call__ run the forward path and trigger variable creation this will trigger the metric variable creation build optimizer note that we mark the state and data to be donated to jax so that jax will reuse the memory buffer for outputs this will reduce the memory usage of the training function by half note that we mark the state and data to be donated to jax so that jax will reuse the memory buffer for outputs this will reduce the memory usage of the training function by half 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 build the model on one batch of data container that configures and calls callbacks callbacks train step setting _jax_state enables callbacks to force a state sync if they need to callbacks reattach state to model variables note doing this after each step would be a big performance bottleneck override with model metrics instead of last step logs the jax spmd_mode is need for multi process context since the metrics values are replicated and we don t want to do a all gather and only need the local copy of the value 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 build the model on one batch of data container that configures and calls callbacks setting _jax_state enables callbacks to force a state sync if they need to i wouldn t recommend modifying non trainable model state during evaluate but it s allowed reattach state back to model the jax spmd_mode is need for multi process context since the metrics values are replicated and we don t want to do a all gather and only need the local copy of the value create an iterator that yields batches of input data build the model on one batch of data build model container that configures and calls callbacks maybe build model train step state sync format return values maybe build model test step state sync format return values build model enforce the sharding spec constraint for all the training state since the output of the train eval step will be used as inputs to next step we need to ensure that they have the same sharding spec so that jax jit won t have to recompile the train eval function note that this function will also rely on the recorded sharding spec for each of states this function is expected to be called within the jitted train eval function especially around the end of the function remove all the model variable for memory saving during jax training since the training function are stateless we have to pass in and get the model weights over and over during which the copy of the weights that attached to the kerasvariable are still and occupying extra memory we remove those variable to save memory for better memory utilization at the beginning of the epoch and reattach the value back to variables at the end of the epoch via jax_state_sync
from functools import partial

import jax
import numpy as np
import tree

from keras import backend
from keras import callbacks as callbacks_module
from keras import ops
from keras import optimizers as optimizers_module
from keras.backend import distribution_lib as jax_distribution_lib
from keras.distribution import distribution_lib
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 JAXTrainer(base_trainer.Trainer):
    def __init__(self):
        super().__init__()
        self.train_function = None
        self.test_function = None
        self.predict_function = None

    def compute_loss_and_updates(
        self,
        trainable_variables,
        non_trainable_variables,
        x,
        y,
        sample_weight,
        training=False,
        optimizer_variables=None,
    ):
        
        kwargs = {}
        if self._call_has_training_arg:
            kwargs["training"] = training
        y_pred, non_trainable_variables, losses = self.stateless_call(
            trainable_variables,
            non_trainable_variables,
            x,
            return_losses=True,
            **kwargs,
        )

        trainable_mapping = zip(self.trainable_variables, trainable_variables)
        with backend.StatelessScope(state_mapping=trainable_mapping):
            
            
            loss = self.compute_loss(
                x, y, y_pred, sample_weight, allow_empty=True
            )
        if losses:
            loss += ops.sum(losses)
        unscaled_loss = loss
        if training and self.optimizer is not None:
            
            mapping = list(zip(self.optimizer.variables, optimizer_variables))
            with backend.StatelessScope(state_mapping=mapping):
                loss = self.optimizer.scale_loss(loss)
        return loss, (unscaled_loss, y_pred, non_trainable_variables)

    def _eager_build(self, data_batch):
        model_unbuilt = not all(layer.built for layer in self._flatten_layers())
        compile_metrics_unbuilt = (
            self._compile_metrics is not None
            and not self._compile_metrics.built
        )
        if model_unbuilt or compile_metrics_unbuilt:

            def _convert_data_to_spec(d):
                if d is None:
                    return None
                return backend.KerasTensor(d.shape, d.dtype)

            data_spec = tree.map_structure(_convert_data_to_spec, data_batch)
            (
                x_spec,
                y_spec,
                sample_weight_spec,
            ) = data_adapter_utils.unpack_x_y_sample_weight(data_spec)
            
            
            y_pred_spec = backend.compute_output_spec(self.__call__, x_spec)
            if compile_metrics_unbuilt:
                
                backend.compute_output_spec(
                    self.compute_metrics,
                    x_spec,
                    y_spec,
                    y_pred_spec,
                    sample_weight=sample_weight_spec,
                )

        if self.optimizer is not None and not self.optimizer.built:
            
            self.optimizer.build(self.trainable_variables)

    def train_step(self, state, data):
        (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        ) = state
        x, y, sample_weight = data_adapter_utils.unpack_x_y_sample_weight(data)
        grad_fn = jax.value_and_grad(
            self.compute_loss_and_updates, has_aux=True
        )
        (loss, aux), grads = grad_fn(
            trainable_variables,
            non_trainable_variables,
            x,
            y,
            sample_weight,
            training=True,
            optimizer_variables=optimizer_variables,
        )
        (unscaled_loss, y_pred, non_trainable_variables) = aux

        (
            trainable_variables,
            optimizer_variables,
        ) = self.optimizer.stateless_apply(
            optimizer_variables, grads, trainable_variables
        )

        with backend.StatelessScope(
            state_mapping=[
                (ref_v, v)
                for ref_v, v in zip(self.metrics_variables, metrics_variables)
            ]
        ) as scope:
            self._loss_tracker.update_state(unscaled_loss)
            logs = self.compute_metrics(x, y, y_pred, sample_weight)

        new_metrics_variables = []
        for ref_v in self.metrics_variables:
            new_v = scope.get_current_value(ref_v)
            if new_v is None:
                new_v = ref_v.value
            new_metrics_variables.append(new_v)
        metrics_variables = new_metrics_variables

        state = self._enforce_jax_state_sharding(
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        )
        return logs, state

    def test_step(self, state, data):
        (
            trainable_variables,
            non_trainable_variables,
            metrics_variables,
        ) = state
        x, y, sample_weight = data_adapter_utils.unpack_x_y_sample_weight(data)
        loss, aux = self.compute_loss_and_updates(
            trainable_variables,
            non_trainable_variables,
            x,
            y,
            sample_weight,
            training=False,
        )
        (unscaled_loss, y_pred, non_trainable_variables) = aux

        with backend.StatelessScope(
            state_mapping=[
                (ref_v, v)
                for ref_v, v in zip(self.metrics_variables, metrics_variables)
            ]
        ) as scope:
            self._loss_tracker.update_state(unscaled_loss)
            logs = self.compute_metrics(x, y, y_pred, sample_weight)

        new_metrics_variables = []
        for ref_v in self.metrics_variables:
            new_v = scope.get_current_value(ref_v)
            if new_v is None:
                new_v = ref_v.value
            new_metrics_variables.append(new_v)
        metrics_variables = new_metrics_variables

        (
            trainable_variables,
            non_trainable_variables,
            _,
            metrics_variables,
        ) = self._enforce_jax_state_sharding(
            trainable_variables=trainable_variables,
            non_trainable_variables=non_trainable_variables,
            optimizer_variables=None,
            metrics_variables=metrics_variables,
        )
        state = (
            trainable_variables,
            non_trainable_variables,
            metrics_variables,
        )
        return logs, state

    def predict_step(self, state, data):
        trainable_variables, non_trainable_variables = state
        kwargs = {}
        if self._call_has_training_arg:
            kwargs["training"] = False

        x, _, _ = data_adapter_utils.unpack_x_y_sample_weight(data)
        outputs, non_trainable_variables = self.stateless_call(
            trainable_variables, non_trainable_variables, x, **kwargs
        )
        (
            trainable_variables,
            non_trainable_variables,
            _,
            _,
        ) = self._enforce_jax_state_sharding(
            trainable_variables=trainable_variables,
            non_trainable_variables=non_trainable_variables,
            optimizer_variables=None,
            metrics_variables=None,
        )
        return outputs, (trainable_variables, non_trainable_variables)

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

        def one_train_step(state, data):
            data = data[0]
            return self.train_step(state, data)

        def multi_train_steps(state, data):
            for single_step_data in data:
                logs, state = one_train_step(state, [single_step_data])
            return logs, state

        if self.steps_per_execution > 1:
            train_step = multi_train_steps
        else:
            train_step = one_train_step

        if not self.run_eagerly and self.jit_compile:
            
            
            
            
            @partial(jax.jit, donate_argnames="state")
            def compiled_train_step(state, data):
                return train_step(state, data)

            self.train_function = compiled_train_step

        else:
            self.train_function = train_step

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

        def one_test_step(state, data):
            data = data[0]
            return self.test_step(state, data)

        def multi_test_steps(state, data):
            for single_step_data in data:
                logs, state = one_test_step(state, [single_step_data])
            return logs, state

        if self.steps_per_execution > 1:
            test_step = multi_test_steps
        else:
            test_step = one_test_step

        if not self.run_eagerly and self.jit_compile:
            
            
            
            
            @partial(jax.jit, donate_argnames="state")
            def compiled_test_step(state, data):
                return test_step(state, data)

            self.test_function = compiled_test_step

        else:
            self.test_function = test_step

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

        def one_predict_step(state, data):
            data = data[0]
            return self.predict_step(state, data)

        def multi_predict_steps(state, data):
            outputs, state = one_predict_step(state, data[:1])

            for single_step_data in data[1:]:
                step_outputs, state = one_predict_step(
                    state,
                    [single_step_data],
                )
                outputs = tree.map_structure(
                    lambda t1, t2: jax.numpy.concatenate([t1, t2]),
                    outputs,
                    step_outputs,
                )
            return outputs, state

        if self.steps_per_execution > 1:
            predict_step = multi_predict_steps
        else:
            predict_step = one_predict_step

        if not self.run_eagerly and self.jit_compile:

            @jax.jit
            def compiled_predict_step(state, data):
                return predict_step(state, data)

            self.predict_function = compiled_predict_step

        else:
            self.predict_function = predict_step

    @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 = EpochIterator(
            x=x,
            y=y,
            sample_weight=sample_weight,
            batch_size=batch_size,
            steps_per_epoch=steps_per_epoch,
            shuffle=shuffle,
            class_weight=class_weight,
            steps_per_execution=self.steps_per_execution,
        )

        needs_building = (
            not all(layer.built for layer in self._flatten_layers())
            or not self.optimizer.built
            or (
                self._compile_metrics is not None
                and not self._compile_metrics.built
            )
        )
        if needs_building:
            
            for _, data in epoch_iterator.enumerate_epoch(return_type="np"):
                data_batch = data[0]
                self._eager_build(data_batch)
                break

        
        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._record_training_state_sharding_spec()

        self.make_train_function()
        self.stop_training = False
        callbacks.on_train_begin()

        for epoch in range(initial_epoch, epochs):
            self.reset_metrics()
            callbacks.on_epoch_begin(epoch)

            trainable_variables = [v.value for v in self.trainable_variables]
            non_trainable_variables = [
                v.value for v in self.non_trainable_variables
            ]
            optimizer_variables = [v.value for v in self.optimizer.variables]
            metrics_variables = [v.value for v in self.metrics_variables]

            self._purge_model_variables()
            for step, data in epoch_iterator.enumerate_epoch(return_type="np"):
                
                callbacks.on_train_batch_begin(step)

                
                state = (
                    trainable_variables,
                    non_trainable_variables,
                    optimizer_variables,
                    metrics_variables,
                )
                data = self._distribute_data(data)
                logs, state = self.train_function(state, data)
                (
                    trainable_variables,
                    non_trainable_variables,
                    optimizer_variables,
                    metrics_variables,
                ) = state

                
                
                self._jax_state = {
                    "trainable_variables": trainable_variables,
                    "non_trainable_variables": non_trainable_variables,
                    "optimizer_variables": optimizer_variables,
                    "metrics_variables": metrics_variables,
                }

                
                callbacks.on_train_batch_end(step, self._pythonify_logs(logs))
                if self.stop_training:
                    break

            
            
            
            self.jax_state_sync()

            
            
            
            
            with jax.spmd_mode("allow_all"):
                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 = EpochIterator(
                        x=val_x,
                        y=val_y,
                        sample_weight=val_sample_weight,
                        batch_size=validation_batch_size or batch_size,
                        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)
        self._jax_state = None
        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 = EpochIterator(
                x=x,
                y=y,
                sample_weight=sample_weight,
                batch_size=batch_size,
                steps_per_epoch=steps,
                shuffle=False,
                steps_per_execution=self.steps_per_execution,
            )

        needs_building = not all(
            layer.built for layer in self._flatten_layers()
        ) or (
            self._compile_metrics is not None
            and not self._compile_metrics.built
        )
        if needs_building:
            
            for _, data in epoch_iterator.enumerate_epoch(return_type="np"):
                data_batch = data[0]
                self._eager_build(data_batch)
                break

        
        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._record_training_state_sharding_spec()

        self.make_test_function()
        self.stop_evaluating = False
        callbacks.on_test_begin()
        logs = None
        self.reset_metrics()

        trainable_variables = [v.value for v in self.trainable_variables]
        non_trainable_variables = [
            v.value for v in self.non_trainable_variables
        ]
        metrics_variables = [v.value for v in self.metrics_variables]

        self._purge_model_variables(optimizer_variables=False)
        for step, data in epoch_iterator.enumerate_epoch(return_type="np"):
            callbacks.on_test_batch_begin(step)

            state = (
                trainable_variables,
                non_trainable_variables,
                metrics_variables,
            )
            data = self._distribute_data(data)
            logs, state = self.test_function(state, data)
            (
                trainable_variables,
                non_trainable_variables,
                metrics_variables,
            ) = state

            
            
            self._jax_state = {
                
                
                "trainable_variables": trainable_variables,
                "non_trainable_variables": non_trainable_variables,
                "metrics_variables": metrics_variables,
            }
            callbacks.on_test_batch_end(step, self._pythonify_logs(logs))
            if self.stop_evaluating:
                break

        
        self.jax_state_sync()

        
        
        
        with jax.spmd_mode("allow_all"):
            logs = self.get_metrics_result()
        callbacks.on_test_end(logs)
        self._jax_state = None
        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 = EpochIterator(
            x=x,
            batch_size=batch_size,
            steps_per_epoch=steps,
            shuffle=False,
            steps_per_execution=self.steps_per_execution,
        )

        if not all(layer.built for layer in self._flatten_layers()):
            
            for _, data in epoch_iterator.enumerate_epoch(return_type="np"):
                
                x, _, _ = data_adapter_utils.unpack_x_y_sample_weight(data[0])
                with backend.StatelessScope():
                    self(x)
                break

        
        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._record_training_state_sharding_spec()

        self.make_predict_function()
        self.stop_predicting = False
        callbacks.on_predict_begin()

        def append_to_outputs(batch_outputs, outputs):
            if outputs is None:
                outputs = tree.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

        trainable_variables = [v.value for v in self.trainable_variables]
        non_trainable_variables = [
            v.value for v in self.non_trainable_variables
        ]
        state = (trainable_variables, non_trainable_variables)
        outputs = None
        for step, x in epoch_iterator.enumerate_epoch(return_type="np"):
            callbacks.on_predict_batch_begin(step)
            x = self._distribute_data(x)
            batch_outputs, state = self.predict_function(state, x)
            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()
        return tree.map_structure_up_to(batch_outputs, np.concatenate, 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")
        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
            )
        data = (x, y, sample_weight)
        data = self._distribute_data(data)

        
        self._eager_build(data)
        self._record_training_state_sharding_spec()
        self.make_train_function()

        
        trainable_variables = [v.value for v in self.trainable_variables]
        non_trainable_variables = [
            v.value for v in self.non_trainable_variables
        ]
        optimizer_variables = [v.value for v in self.optimizer.variables]
        metrics_variables = [v.value for v in self.metrics_variables]
        state = (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        )
        logs, state = self.train_function(state, [data])

        
        (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        ) = state
        self._jax_state = {
            "trainable_variables": trainable_variables,
            "non_trainable_variables": non_trainable_variables,
            "optimizer_variables": optimizer_variables,
            "metrics_variables": metrics_variables,
        }
        self.jax_state_sync()

        
        logs = tree.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")

        data = (x, y, sample_weight)
        data = self._distribute_data(data)
        
        self._eager_build(data)
        self._record_training_state_sharding_spec()
        self.make_test_function()

        
        trainable_variables = [v.value for v in self.trainable_variables]
        non_trainable_variables = [
            v.value for v in self.non_trainable_variables
        ]
        metrics_variables = [v.value for v in self.metrics_variables]
        state = (
            trainable_variables,
            non_trainable_variables,
            metrics_variables,
        )
        logs, state = self.test_function(state, [data])

        
        trainable_variables, non_trainable_variables, metrics_variables = state
        self._jax_state = {
            "trainable_variables": trainable_variables,
            "non_trainable_variables": non_trainable_variables,
            "metrics_variables": metrics_variables,
        }
        self.jax_state_sync()

        
        logs = tree.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):
        if not all(layer.built for layer in self._flatten_layers()):
            
            with backend.StatelessScope():
                self(x)
        self._record_training_state_sharding_spec()
        self.make_predict_function()

        trainable_variables = [v.value for v in self.trainable_variables]
        non_trainable_variables = [
            v.value for v in self.non_trainable_variables
        ]
        state = (trainable_variables, non_trainable_variables)
        batch_outputs, state = self.predict_function(state, [(x,)])
        batch_outputs = tree.map_structure(lambda x: np.array(x), batch_outputs)
        return batch_outputs

    def jax_state_sync(self):
        if not getattr(self, "_jax_state", None):
            return

        trainable_variables = self._jax_state.get("trainable_variables", None)
        non_trainable_variables = self._jax_state.get(
            "non_trainable_variables", None
        )
        optimizer_variables = self._jax_state.get("optimizer_variables", None)
        metrics_variables = self._jax_state.get("metrics_variables", None)
        if trainable_variables:
            for ref_v, v in zip(self.trainable_variables, trainable_variables):
                ref_v.assign(v)
        if non_trainable_variables:
            for ref_v, v in zip(
                self.non_trainable_variables, non_trainable_variables
            ):
                ref_v.assign(v)
        if optimizer_variables:
            for ref_v, v in zip(self.optimizer.variables, optimizer_variables):
                ref_v.assign(v)
        if metrics_variables:
            for ref_v, v in zip(self.metrics_variables, metrics_variables):
                ref_v.assign(v)

    def _distribute_data(self, data):
        distribution = distribution_lib.distribution()
        if distribution is not None:

            def distribute_single_value(d):
                layout = distribution.get_data_layout(d.shape)
                return jax_distribution_lib.distribute_data_input(d, layout)

            return jax.tree_util.tree_map(distribute_single_value, data)
        else:
            return data

    def _record_training_state_sharding_spec(self):
        self._trainable_variable_shardings = [
            v.value.sharding for v in self.trainable_variables
        ]
        self._non_trainable_variable_shardings = [
            v.value.sharding for v in self.non_trainable_variables
        ]
        if hasattr(self, "optimizer") and self.optimizer is not None:
            self._optimizer_variable_shardings = [
                v.value.sharding for v in self.optimizer.variables
            ]
        else:
            self._optimizer_variable_shardings = []
        self._metrics_variable_shardings = [
            v.value.sharding for v in self.metrics_variables
        ]

    def _enforce_jax_state_sharding(
        self,
        trainable_variables=None,
        non_trainable_variables=None,
        optimizer_variables=None,
        metrics_variables=None,
    ):
        
        trainable_variables = trainable_variables or []
        non_trainable_variables = non_trainable_variables or []
        optimizer_variables = optimizer_variables or []
        metrics_variables = metrics_variables or []

        for i in range(len(trainable_variables)):
            trainable_variables[i] = jax.lax.with_sharding_constraint(
                trainable_variables[i], self._trainable_variable_shardings[i]
            )
        for i in range(len(non_trainable_variables)):
            non_trainable_variables[i] = jax.lax.with_sharding_constraint(
                non_trainable_variables[i],
                self._non_trainable_variable_shardings[i],
            )
        for i in range(len(optimizer_variables)):
            optimizer_variables[i] = jax.lax.with_sharding_constraint(
                optimizer_variables[i], self._optimizer_variable_shardings[i]
            )
        for i in range(len(metrics_variables)):
            metrics_variables[i] = jax.lax.with_sharding_constraint(
                metrics_variables[i], self._metrics_variable_shardings[i]
            )
        return (
            trainable_variables,
            non_trainable_variables,
            optimizer_variables,
            metrics_variables,
        )

    def _purge_model_variables(
        self,
        trainable_variables=True,
        non_trainable_variables=True,
        optimizer_variables=True,
        metric_variables=True,
    ):
        
        if trainable_variables:
            for v in self.trainable_variables:
                v._value = None
        if non_trainable_variables:
            for v in self.non_trainable_variables:
                v._value = None
        if optimizer_variables:
            for v in self.optimizer.variables:
                v._value = None
        if metric_variables:
            for v in self.metrics_variables:
                v._value = None