keras/src/optimizers/base_optimizer.py/BaseOptimizer/_update_model_variables_moving_average
class BaseOptimizer:    def _update_model_variables_moving_average(self, trainable_variables):
        """Update the stored moving average using the latest value."""
        if self.use_ema:
            for var, average in zip(
                trainable_variables, self._model_variables_moving_average
            ):
                not_first_step = ops.not_equal(self.iterations, 0)
                momentum = (
                    ops.cast(not_first_step, var.dtype) * self.ema_momentum
                )
                average.assign(momentum * average + (1 - momentum) * var)