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)
        local_step = ops.cast(self.iterations + 1, variable.dtype)

        beta_1_power = ops.power(
            ops.cast(self.beta_1, variable.dtype), local_step
        )
        beta_2_power = ops.power(
            ops.cast(self.beta_2, variable.dtype), local_step
        )

        m = self._momentums[self._get_variable_index(variable)]
        v = self._velocities[self._get_variable_index(variable)]

        self.assign_add(
            m, ops.multiply(ops.subtract(gradient, m), 1 - self.beta_1)
        )

        self.assign_add(
            v,
            ops.multiply(
                ops.subtract(ops.square(gradient), v), 1 - self.beta_2
            ),
        )

        m_t_hat = ops.divide(m, (1.0 - beta_1_power))
        v_sqrt = ops.add(
            ops.sqrt(ops.divide(v, (1.0 - beta_2_power))), self.epsilon
        )

        update = ops.divide(m_t_hat, v_sqrt)
        w_norm = ops.sqrt(ops.sum(ops.power(variable, 2)))
        g_norm = ops.sqrt(ops.sum(ops.power(update, 2)))

        # ratio = w_norm / g_norm if w_norm > 0 and g_norm > 0 else 1
        ratio = ops.where(
            ops.greater(w_norm, 0),
            ops.where(ops.greater(g_norm, 0), (w_norm / g_norm), 1.0),
            1.0,
        )

        self.assign_sub(variable, ratio * lr * update)
Update step given gradient and the associated model variable.
keras/src/optimizers/lamb.py