def _backend_apply_gradients(self, grads, trainable_variables):
        """Apply method that can be overridden by different backends.

        JAX overrides it in order to deal with statelessness in gradient
        accumulation and EMA handling.

        The below implementation is intended to be generally backend-agnostic,
        but may not work with all backends.

        This method does 4 things:
        - Call the optimizer's update_step() to update trainable variables
            and optimizer variables.
        - Update EMA variables, if EMA is configured.
        - Update gradient accumulators, if gradient accumulation is configured.
        - Update the iteration counter.
        """
        if self.gradient_accumulation_steps:
            is_update_step = (
                self._iterations + 1
            ) % self.gradient_accumulation_steps == 0
            # `trainable_variables` might have been filtered in previous
            # processing steps, so we need to ensure the correct mapping between
            # `self._accumulated_gradients` and `trainable_variables`
            acc_grads = [
                self._accumulated_gradients[self._get_variable_index(v)]
                for v in trainable_variables
            ]

            def _update_step_fn(grads, trainable_variables):
                # Run update step with accumulated grads + reset accumulators
                steps = self.gradient_accumulation_steps
                grads = [
                    (g + acc_g) / steps for g, acc_g in zip(grads, acc_grads)
                ]

                # Apply clipping and weight decay.
                grads = self._clip_gradients(grads)
                self._apply_weight_decay(trainable_variables)

                self._backend_update_step(
                    grads, trainable_variables, self.learning_rate
                )
                self._backend_reset_gradient_accumulators()

            ops.cond(
                is_update_step,
                lambda: _update_step_fn(grads, trainable_variables),
                lambda: self._backend_increment_gradient_accumulators(
                    grads, acc_grads
                ),
            )
        else:
            # Apply clipping and weight decay.
            grads = self._clip_gradients(grads)
            self._apply_weight_decay(trainable_variables)

            # Run update step.
            self._backend_update_step(
                grads, trainable_variables, self.learning_rate
            )

        if self.use_ema:
            self._update_model_variables_moving_average(
                self._trainable_variables
            )
            if self.ema_overwrite_frequency:
                # Only when self.ema_overwrite_frequency is not None, we
                # overwrite the model variables.
                should_overwrite_model_vars = (
                    self.iterations + 1
                ) % self.ema_overwrite_frequency == 0
                ops.cond(
                    should_overwrite_model_vars,
                    lambda: self._overwrite_model_variables_with_average_value(
                        self._trainable_variables
                    ),
                    lambda: None,
                )
Apply method that can be overridden by different backends.

        JAX overrides it in order to deal with statelessness in gradient
        accumulation and EMA handling.

        The below implementation is intended to be generally backend-agnostic,
        but may not work with all backends.

        This method does 4 things:
        - Call the optimizer's update_step() to update trainable variables
            and optimizer variables.
        - Update EMA variables, if EMA is configured.
        - Update gradient accumulators, if gradient accumulation is configured.
        - Update the iteration counter.
        
keras/src/optimizers/base_optimizer.py