def METHOD_NAME(decay_steps, num_periods=1, zero_after=None):
  """Returns a function that computes a restart decay.
  This decay computes
    0.5 * (1.0 + cos(pi * (num_periods * global_step) % num_training_steps))
  This is a simplified version of the restart decay introduced in
  "SGDR: Stochastic Gradient Descent with Warm Restarts"
  by Ilya Loshchilov & Frank Hutter, Proceedings of
  ICLR'2017, available at https://arxiv.org/pdf/1608.03983.pdf
  This decay can be used to decay the sign quantity in the AddSign and PowerSign
  optimizers discovered in
  [Bello et al., ICML 2017] Neural Optimizer Search with RL.
  Example usage:
  ```
  decay_steps = 1000
  num_periods = 2.0
  restart_decay_fn = get_restart_decay_fn(decay_steps,
                                          num_periods=num_periods)
  decayed = restart_decay_fn(global_step)
  x *= decayed
  ```
  Args:
    decay_steps: number of steps to decay over.
    num_periods: number of periods for cosine signal. 1 by default,
      which maps the last decay step to 0.
    zero_after: if not None, number after which the decay function
      will return 0.
  Returns:
    restart_decay_fn: a function that computes the restart decay.
  """
  # pylint:disable=missing-docstring
  def restart_decay_fn(global_step):
    if global_step is None:
      raise ValueError("global_step is required for cosine_decay.")
    global_step = math_ops.minimum(global_step, decay_steps)
    num = math_ops.mod(num_periods * math_ops.cast(global_step, dtypes.float32),
                       decay_steps)
    fraction = num / math_ops.cast(decay_steps, dtypes.float32)
    decayed = 0.5 * (
        1.0 + math_ops.cos(constant_op.constant(math.pi) * fraction))
    if zero_after is not None:
      tmp = (math_ops.cast(num_periods * global_step, dtypes.float32) /
             math_ops.cast(decay_steps, dtypes.float32))
      decayed = array_ops.where(
          math_ops.greater_equal(tmp, zero_after), 0.0, decayed)
    return decayed
  # pylint:enable=missing-docstring
  return restart_decay_fn