keras/src/optimizers/schedules/learning_rate_schedule.py/CosineDecayRestarts/compute_step
class CosineDecayRestarts:    def __call__(self, step):
        with ops.name_scope(self.name):
            initial_learning_rate = ops.convert_to_tensor(
                self.initial_learning_rate
            )
            dtype = initial_learning_rate.dtype
            first_decay_steps = ops.cast(self.first_decay_steps, dtype)
            alpha = ops.cast(self.alpha, dtype)
            t_mul = ops.cast(self._t_mul, dtype)
            m_mul = ops.cast(self._m_mul, dtype)

            global_step_recomp = ops.cast(step, dtype)
            completed_fraction = global_step_recomp / first_decay_steps

            def compute_step(completed_fraction, geometric=False):
                """Helper for `cond` operation."""
                if geometric:
                    # ops.log is sensitive to the precision of dtype, so we need
                    # the additional casting
                    i_restart = ops.floor(
                        ops.log(
                            ops.cast(
                                1.0 - completed_fraction * (1.0 - t_mul), dtype
                            )
                        )
                        / ops.log(t_mul)
                    )

                    sum_r = (1.0 - t_mul**i_restart) / (1.0 - t_mul)
                    completed_fraction = (
                        completed_fraction - sum_r
                    ) / t_mul**i_restart

                else:
                    i_restart = ops.floor(completed_fraction)
                    completed_fraction -= i_restart

                return i_restart, completed_fraction

            i_restart, completed_fraction = ops.cond(
                ops.equal(t_mul, 1.0),
                lambda: compute_step(completed_fraction, geometric=False),
                lambda: compute_step(completed_fraction, geometric=True),
            )

            m_fac = m_mul**i_restart
            cosine_decayed = (
                0.5
                * m_fac
                * (
                    1.0
                    + ops.cos(
                        ops.array(math.pi, dtype=dtype) * completed_fraction
                    )
                )
            )
            decayed = (1 - alpha) * cosine_decayed + alpha

            return ops.multiply(initial_learning_rate, decayed)