def _overwrite_variables_directly_with_gradients(self, grads, vars):
        """Overwrite the variables directly by their gradients.

        This method is designed for a special case where we want to overwrite
        the variable directly with its computed gradient. For example, in float8
        training, new `scale` and `amax_history` are computed as gradients, and
        we want to overwrite them directly instead of following the typical
        procedure such as gradient descent with a learning rate, gradient
        clipping and weight decaying.

        After the update, the processed pairs will be filtered out.
        """
        # Shortcut for `tf.Variable` because it doesn't have a
        # `overwrite_with_gradient` attr.
        if not any(self._overwrite_variable_with_gradient(v) for v in vars):
            return grads, vars

        # Shallow copies
        filtered_grads = list(grads)
        filtered_vars = list(vars)

        # Iterate from right to left for safe popping
        for i in range(len(filtered_grads) - 1, -1, -1):
            g, v = filtered_grads[i], filtered_vars[i]
            if self._overwrite_variable_with_gradient(v):
                if self.gradient_accumulation_steps:
                    # Utilize a stateless manner for JAX compatibility
                    steps = self.gradient_accumulation_steps
                    is_update_step = (self._iterations + 1) % steps == 0
                    acc_g = self._accumulated_gradients[
                        self._get_variable_index(v)
                    ]
                    # `ops.maximum` is utilized for gradient accumulation for
                    # `overwrite_with_gradient=True` variables
                    new_g_acc = ops.cond(
                        is_update_step,
                        lambda: ops.zeros(g.shape, dtype=g.dtype),
                        lambda: ops.maximum(g, acc_g),
                    )
                    new_g = ops.cond(
                        is_update_step,
                        lambda: ops.maximum(g, acc_g),
                        lambda: g,
                    )
                    new_v = ops.cond(
                        is_update_step, lambda: new_g, lambda: v.value
                    )
                    v.assign(new_v)
                    acc_g.assign(new_g_acc)
                else:
                    v.assign(g)
                filtered_grads.pop(i)
                filtered_vars.pop(i)
        return filtered_grads, filtered_vars
Overwrite the variables directly by their gradients.

        This method is designed for a special case where we want to overwrite
        the variable directly with its computed gradient. For example, in float8
        training, new `scale` and `amax_history` are computed as gradients, and
        we want to overwrite them directly instead of following the typical
        procedure such as gradient descent with a learning rate, gradient
        clipping and weight decaying.

        After the update, the processed pairs will be filtered out.
        
keras/src/optimizers/base_optimizer.py