code: def create_learning_rate_scheduler(
      factors="constant * linear_warmup * rsqrt_decay",
      base_learning_rate=0.5,
      warmup_steps=8000,
      decay_factor=0.5,
      steps_per_decay=20000,
      steps_per_cycle=100000,
  ):
      """creates 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_decay: Cyclic cosine decay, uses steps_per_cycle parameter.
  
      Args:
        factors: a string with factors separated by '*' that defines the schedule.
        base_learning_rate: float, the starting constant for the lr 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.
  
      Returns:
        a function learning_rate(step): float -> {'learning_rate': float}, the
        step-dependent lr.
      """
      factors = [n.strip() for n in factors.split("*")]
  
      def step_fn(step):
          """Step to learning rate function."""
          ret = 1.0
          for name in factors:
              if name == "constant":
                  ret *= base_learning_rate
              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 *= jnp.maximum(
                      0.0,
                      0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))),
                  )
              else:
                  raise ValueError("Unknown factor %s." % name)
          return jnp.asarray(ret, dtype=jnp.float32)
  
      return step_fn
  
imports: import functools
  import os
  import time
  from absl import app
  from absl import flags
  from absl import logging
  from flax import jax_utils
  from flax import linen
  from flax.metrics import tensorboard
  from flax.training import common_utils
  from flax.training import train_state
  import jax
  import jax.numpy
  from jax import random
  import numpy
  import optax
  import tensorflow
  import input_pipeline
  import models
label_desc: Use image node_id 1 for calling a global function with example usage: create_learning_rate_scheduler(factors, base_learning_rate, warmup_steps, decay_factor, steps_per_decay, steps_per_cycle) and returns: step_fn, jnp