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)

        accum = self._accumulators[self._get_variable_index(variable)]
        linear = self._linears[self._get_variable_index(variable)]

        lr_power = self.learning_rate_power
        l2_reg = self.l2_regularization_strength
        l2_reg = l2_reg + self.beta / (2.0 * lr)

        grad_to_use = ops.add(
            gradient,
            ops.multiply(
                2 * self.l2_shrinkage_regularization_strength, variable
            ),
        )
        new_accum = ops.add(accum, ops.square(gradient))
        self.assign_add(
            linear,
            ops.subtract(
                grad_to_use,
                ops.multiply(
                    ops.divide(
                        ops.subtract(
                            ops.power(new_accum, -lr_power),
                            ops.power(accum, -lr_power),
                        ),
                        lr,
                    ),
                    variable,
                ),
            ),
        )
        quadratic = ops.add(
            ops.divide(ops.power(new_accum, (-lr_power)), lr), 2 * l2_reg
        )
        linear_clipped = ops.clip(
            linear,
            -self.l1_regularization_strength,
            self.l1_regularization_strength,
        )
        self.assign(
            variable,
            ops.divide(ops.subtract(linear_clipped, linear), quadratic),
        )
        self.assign(accum, new_accum)
Update step given gradient and the associated model variable.
keras/src/optimizers/ftrl.py