keras/src/optimizers/schedules/learning_rate_schedule.py/CosineDecay/__call__
class CosineDecay:    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
            decay_steps = ops.cast(self.decay_steps, dtype)
            global_step_recomp = ops.cast(step, dtype)

            if self.warmup_target is None:
                global_step_recomp = ops.minimum(
                    global_step_recomp, decay_steps
                )
                return self._decay_function(
                    global_step_recomp,
                    decay_steps,
                    initial_learning_rate,
                    dtype,
                )

            warmup_target = ops.cast(self.warmup_target, dtype)
            warmup_steps = ops.cast(self.warmup_steps, dtype)

            global_step_recomp = ops.minimum(
                global_step_recomp, decay_steps + warmup_steps
            )

            return ops.cond(
                global_step_recomp < warmup_steps,
                lambda: self._warmup_function(
                    global_step_recomp,
                    warmup_steps,
                    warmup_target,
                    initial_learning_rate,
                ),
                lambda: self._decay_function(
                    global_step_recomp - warmup_steps,
                    decay_steps,
                    warmup_target,
                    dtype,
                ),
            )