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)
        epsilon_2 = ops.cast(self.epsilon_2, variable.dtype)
        one = ops.cast(1.0, variable.dtype)
        local_step = ops.cast(self.iterations + 1, variable.dtype)
        if not callable(self._learning_rate) and self.relative_step:
            lr = ops.minimum(lr, 1 / ops.sqrt(local_step))

        r = self._r[self._get_variable_index(variable)]
        c = self._c[self._get_variable_index(variable)]
        v = self._v[self._get_variable_index(variable)]

        rho_t = ops.minimum(lr, 1 / ops.sqrt(local_step))
        alpha_t = ops.maximum(epsilon_2, self._rms(variable)) * rho_t
        regulated_grad_square = ops.add(ops.square(gradient), self.epsilon_1)
        beta_2_t = 1 - ops.power(local_step, self.beta_2_decay)

        if len(variable.shape) >= 2:
            # `r` deletes the last dimension of gradient, so it is of shape
            # `gradient.shape[:-1]`.
            self.assign(
                r,
                beta_2_t * r
                + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-1),
            )
            # `c` deletes the second last dimension of gradient, so it is of
            # shape `gradient.shape[:-2] + gradient.shape[-1]`.
            self.assign(
                c,
                beta_2_t * c
                + (1 - beta_2_t) * ops.mean(regulated_grad_square, axis=-2),
            )
            self.assign(
                v,
                ops.expand_dims(
                    r / ops.mean(r, axis=-1, keepdims=True), axis=-1
                )
                * ops.expand_dims(c, -2),
            )
        else:
            self.assign(
                v, beta_2_t * v + (1 - beta_2_t) * regulated_grad_square
            )

        u_t = ops.divide(gradient, ops.sqrt(v))
        u_t_hat = ops.divide(
            u_t,
            ops.maximum(one, ops.divide(self._rms(u_t), self.clip_threshold)),
        )
        self.assign_sub(variable, ops.multiply(alpha_t, u_t_hat))
Update step given gradient and the associated model variable.
keras/src/optimizers/adafactor.py