a class for tensorflow specific optimizer logic the major behavior change for this class is for tf distribute it will override methods from base keras core optimizer which provide distribute specific functionality e g variable creation loss reduction etc this is mainly due to the interaction with tf distribute strategy which requires tf variable as the inputs for most of its apis for resourcevariables the distributedcontainer attribute is added to their handle tensors applygradients using a distributionstrategy unwrappedgradsandvars for g v in gradsandvars if isinstancev backend variable v v value unwrappedgradsandvars appendg v gradsandvars unwrappedgradsandvars def applygradtoupdatevarvar grad learningrate self getcurrentlearningrate return self updatestepgrad var learningrate for grad var in gradsandvars distribution extended update var applygradtoupdatevar argsgrad groupfalse if self useema varlist zipgradsandvars self updatemodelvariablesmovingaveragevarlist if self emaoverwritefrequency only when self emaoverwritefrequency is not none we overwrite the model variables shouldoverwritemodelvars self iterations 1 self emaoverwritefrequency 0 tf cond tf castshouldoverwritemodelvars tf bool truefnlambda self overwritemodelvariableswithaveragevalue noqa e501 varlist falsefnlambda none self iterations assignadd1 def overwritemodelvariableswithaveragevalueself varlist strategy self distributionstrategy override model variable by the stored average value on all devices for var averagevar in zip varlist self modelvariablesmovingaverage strategy extended update var lambda a b a assignb argsaveragevar a class for tensorflow specific optimizer logic the major behavior change for this class is for tf distribute it will override methods from base keras core optimizer which provide distribute specific functionality e g variable creation loss reduction etc this is mainly due to the interaction with tf distribute strategy which requires tf variable as the inputs for most of its apis convert to tf variable for resourcevariables the _distributed_container attribute is added to their handle tensors convert to tf variable apply_gradients using a distributionstrategy only when self ema_overwrite_frequency is not none we overwrite the model variables noqa e501 overwrite model variables with their moving average values this function overwrites variables on each device args var_list list of model variables override model variable by the stored average value on all devices
import tensorflow as tf

from keras import backend
from keras.backend.common import KerasVariable
from keras.optimizers import base_optimizer


class TFOptimizer(base_optimizer.BaseOptimizer):
    

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._distribution_strategy = tf.distribute.get_strategy()

    def add_variable_from_reference(self, reference_variable, name=None):
        if isinstance(reference_variable, backend.Variable):
            colocate_var = reference_variable.value
        else:
            colocate_var = reference_variable

        with self._distribution_strategy.extended.colocate_vars_with(
            colocate_var
        ):
            return super().add_variable_from_reference(
                reference_variable, name=name
            )

    def stateless_apply(self, optimizer_variables, grads, trainable_variables):
        
        
        raise ValueError(
            "stateless_apply is not supported with the TensorFlow backend "
            "(as it is incompatible with tf.distribute)."
        )

    def assign_add(self, variable, value):
        if isinstance(variable, KerasVariable):
            variable = variable.value
        value = tf.cast(value, variable.dtype)
        if isinstance(value, tf.IndexedSlices):
            variable.scatter_add(value)
        else:
            variable.assign_add(value)

    def assign_sub(self, variable, value):
        if isinstance(variable, KerasVariable):
            variable = variable.value
        value = tf.cast(value, variable.dtype)
        if isinstance(value, tf.IndexedSlices):
            variable.scatter_sub(value)
        else:
            variable.assign_sub(value)

    def _var_key(self, variable):
        if isinstance(variable, backend.Variable):
            variable = variable.value  
        if hasattr(variable, "_distributed_container"):
            variable = variable._distributed_container()
        elif (
            isinstance(variable, tf.__internal__.CompositeTensor)
            and hasattr(variable, "handle")
            and hasattr(variable.handle, "_distributed_container")
        ):
            
            
            variable = variable.handle._distributed_container()
        return variable._unique_id

    def _apply_weight_decay(self, variables):
        if self.weight_decay is None:
            return

        def distributed_apply_weight_decay(distribution, variables, **kwargs):
            def weight_decay_fn(variable):
                if self._use_weight_decay(variable):
                    lr = tf.cast(self.learning_rate, variable.dtype)
                    wd = tf.cast(self.weight_decay, variable.dtype)
                    variable.assign_sub(variable * wd * lr)

            for variable in variables:
                if isinstance(variable, backend.Variable):
                    variable = variable.value  
                distribution.extended.update(
                    variable, weight_decay_fn, group=False
                )

        tf.__internal__.distribute.interim.maybe_merge_call(
            distributed_apply_weight_decay,
            self._distribution_strategy,
            variables,
        )

    def _internal_apply_gradients(self, grads_and_vars):
        tf.__internal__.distribute.interim.maybe_merge_call(
            self._distributed_apply_gradients_fn,
            self._distribution_strategy,
            grads_and_vars,
        )

    def _distributed_apply_gradients_fn(
        self, distribution, grads_and_vars, **kwargs
    ):
        

        unwrapped_grads_and_vars = []
        for g, v in grads_and_vars:
            if isinstance(v, backend.Variable):
                v = v.value
            unwrapped_grads_and_vars.append((g, v))
        grads_and_vars = unwrapped_grads_and_vars

        def apply_grad_to_update_var(var, grad):
            learning_rate = self._get_current_learning_rate()
            return self.update_step(grad, var, learning_rate)

        for grad, var in grads_and_vars:
            distribution.extended.update(
                var, apply_grad_to_update_var, args=(grad,), group=False
            )

        if self.use_ema:
            _, var_list = zip(*grads_and_vars)
            self._update_model_variables_moving_average(var_list)
            if self.ema_overwrite_frequency:
                
                
                should_overwrite_model_vars = (
                    self.iterations + 1
                ) % self.ema_overwrite_frequency == 0
                tf.cond(
                    tf.cast(should_overwrite_model_vars, tf.bool),
                    true_fn=lambda: self._overwrite_model_variables_with_average_value(  
                        var_list
                    ),
                    false_fn=lambda: None,
                )
        self.iterations.assign_add(1)

    def _overwrite_model_variables_with_average_value(self, var_list):
        
        strategy = self._distribution_strategy
        
        for var, average_var in zip(
            var_list, self._model_variables_moving_average
        ):
            strategy.extended.update(
                var, lambda a, b: a.assign(b), args=(average_var,)
            )