def multifactor(factors='constant * linear_warmup * rsqrt_decay',
                constant=0.1,  # pylint: disable=redefined-outer-name
                warmup_steps=400,
                decay_factor=0.5,
                steps_per_decay=20000,
                steps_per_cycle=100000,
                second_constant=0.01,
                second_constant_step=10000,
                minimum=0):
  """Factor-based learning rate schedule.

  Interprets factors in the factors string which can consist of:
  * constant: interpreted as the constant value,
  * linear_warmup: interpreted as linear warmup until warmup_steps,
  * rsqrt_decay: divide by square root of max(step, warmup_steps)
  * decay_every: Every k steps decay the learning rate by decay_factor.
  * cosine_deay: Cyclic cosine decay, uses steps_per_cycle parameter.
  * two_constants: constant until second_constant_step, then switch to
    second_constant.

  Args:
    factors: a string with factors separated by '*' that defines the schedule.
    constant: float, the starting constant for the learning rate schedule.
    warmup_steps: how many steps to warm up for in the warmup schedule.
    decay_factor: The amount to decay the learning rate by.
    steps_per_decay: How often to decay the learning rate.
    steps_per_cycle: Steps per cycle when using cosine decay.
    second_constant: float, the second constant for the learning rate schedule.
    second_constant_step: the step when the second_constant is triggered.
    minimum: if the computed rate is below the minimum, then return the minimum.

  Returns:
    a function learning_rate(step): float -> {'learning_rate': float}, the
    step-dependent lr.
  """
  factors = [n.strip() for n in factors.split('*')]

  def learning_rate(step):
    """Step to learning rate function."""
    ret = 1.0
    for name in factors:
      if name == 'constant':
        ret *= constant
      elif name == 'two_constants':
        if step < second_constant_step:
          ret *= constant
        else:
          ret *= second_constant
      elif name == 'linear_warmup':
        ret *= jnp.minimum(1.0, step / warmup_steps)
      elif name == 'rsqrt_decay':
        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
      elif name == 'rsqrt_normalized_decay':
        ret *= jnp.sqrt(warmup_steps)
        ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
      elif name == 'decay_every':
        ret *= (decay_factor ** (step//steps_per_decay))
      elif name == 'cosine_decay':
        progress = jnp.maximum(
            0.0, (step - warmup_steps) / float(steps_per_cycle))
        ret *= (0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))))
      else:
        raise ValueError('Unknown factor %s.' % name)
    # TODO(henrykm): return float(jnp.max(minimum, ret)) would be
    # better but causes TypeError: 'numpy.float64' object cannot
    # be interpreted as an integer
    if ret <= minimum:
      return minimum
    return ret

  return learning_rate
Factor-based learning rate schedule.

  Interprets factors in the factors string which can consist of:
  * constant: interpreted as the constant value,
  * linear_warmup: interpreted as linear warmup until warmup_steps,
  * rsqrt_decay: divide by square root of max(step, warmup_steps)
  * decay_every: Every k steps decay the learning rate by decay_factor.
  * cosine_deay: Cyclic cosine decay, uses steps_per_cycle parameter.
  * two_constants: constant until second_constant_step, then switch to
    second_constant.

  Args:
    factors: a string with factors separated by '*' that defines the schedule.
    constant: float, the starting constant for the learning rate schedule.
    warmup_steps: how many steps to warm up for in the warmup schedule.
    decay_factor: The amount to decay the learning rate by.
    steps_per_decay: How often to decay the learning rate.
    steps_per_cycle: Steps per cycle when using cosine decay.
    second_constant: float, the second constant for the learning rate schedule.
    second_constant_step: the step when the second_constant is triggered.
    minimum: if the computed rate is below the minimum, then return the minimum.

  Returns:
    a function learning_rate(step): float -> {'learning_rate': float}, the
    step-dependent lr.
  
trax/supervised/lr_schedules.py