keras/src/callbacks/swap_ema_weights.py/SwapEMAWeights/_tf_swap_variables
class SwapEMAWeights:    def _tf_swap_variables(self, optimizer):
        for var, average_var in zip(
            self.model.trainable_variables,
            optimizer._model_variables_moving_average,
        ):
            if isinstance(var, backend.Variable):
                var = var.value
            if isinstance(average_var, backend.Variable):
                average_var = average_var.value
            # swap using addition to prevent variable creation
            optimizer._distribution_strategy.extended.update(
                var,
                lambda a, b: a.assign_add(b),
                args=(average_var,),
            )
            optimizer._distribution_strategy.extended.update(
                var,
                lambda a, b: b.assign(a - b),
                args=(average_var,),
            )
            optimizer._distribution_strategy.extended.update(
                var,
                lambda a, b: a.assign(a - b),
                args=(average_var,),
            )