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
            ):
                if average is not None:
                    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)
Update the stored moving average using the latest value.
keras/src/optimizers/base_optimizer.py