code: def create_learning_rate_scheduler(
      factors="constant * linear_warmup * rsqrt_decay",
      base_learning_rate=0.5,
      warmup_steps=1000,
      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)
      * rsqrt_normalized_decay: divide by square root of max(step/warmup_steps, 1)
      * decay_every: Every k steps decay the learning rate by decay_factor.
      * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter.
      Args:
        factors: string, factors separated by "*" that defines the schedule.
        base_learning_rate: float, the starting constant for the lr schedule.
        warmup_steps: int, how many steps to warm up for in the warmup schedule.
        decay_factor: float, the amount to decay the learning rate by.
        steps_per_decay: int, how often to decay the learning rate.
        steps_per_cycle: int, 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 logging
  import os
  import sys
  from dataclasses import dataclass, field
  from pathlib import Path
  from typing import Dict, List, Optional, Tuple
  import jax
  import jax.numpy
  import numpy
  from datasets import load_dataset
  from flax import jax_utils
  from flax.optim import Adam
  from flax.training import common_utils
  from flax.training.common_utils import get_metrics
  from jax.nn import log_softmax
  from modeling_flax_performer import FlaxPerformerForMaskedLM
  from tqdm import tqdm
  from transformers import MODEL_FOR_MASKED_LM_MAPPING, AutoTokenizer, BertConfig, FlaxBertForMaskedLM, HfArgumentParser, PreTrainedTokenizerBase, TensorType, TrainingArguments, is_tensorboard_available, set_seed
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