def update_step(self, gradient, variable, learning_rate):
        """Update step given gradient and the associated model variable."""
        lr = ops.cast(learning_rate, variable.dtype)
        gradient = ops.cast(gradient, variable.dtype)

        velocity = self._velocities[self._get_variable_index(variable)]
        momentum = None
        if self.momentum > 0:
            momentum = self._momentums[self._get_variable_index(variable)]
        average_grad = None
        if self.centered:
            average_grad = self._average_gradients[
                self._get_variable_index(variable)
            ]

        rho = self.rho

        self.assign(
            velocity,
            ops.add(
                ops.multiply(rho, velocity),
                ops.multiply(1 - rho, ops.square(gradient)),
            ),
        )
        if self.centered:
            self.assign(
                average_grad,
                ops.add(
                    ops.multiply(rho, average_grad),
                    ops.multiply(1 - rho, gradient),
                ),
            )
            denominator = velocity - ops.square(average_grad) + self.epsilon
        else:
            denominator = ops.add(velocity, self.epsilon)
        increment = ops.divide(
            ops.multiply(lr, gradient), ops.sqrt(denominator)
        )
        if self.momentum > 0:
            self.assign(
                momentum,
                ops.add(ops.multiply(self.momentum, momentum), increment),
            )
            self.assign_sub(variable, momentum)
        else:
            self.assign_sub(variable, increment)
Update step given gradient and the associated model variable.
keras/src/optimizers/rmsprop.py