code: 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
  
file_path: trax/trax/supervised/lr_schedules.py
imports: import math
  import gin
  from trax.fastmath import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: multifactor(factors, constant, warmup_steps, decay_factor, steps_per_decay, steps_per_cycle, second_constant, second_constant_step, minimum) and returns: learning_rate, ret, minimum