code: def rsqrt_schedule(
      init_value: float,
      shift: int = 0,
  ):
      """Applies a reverse square-root schedule.
  
      The reverse square root schedule is simply `lr = init_value / sqrt(step)`.
  
      Args:
        init_value: Base learning rate (before applying the rsqrt schedule).
        shift: How many steps the rsqrt should be shifted. Shifting the rsqrt
          schedule makes it less steep in the beginning (close to 0).
  
      Returns:
        A schedule that applies the reverse square root.
      """
  
      def schedule(count):
          return init_value * (count + shift) ** -0.5 * shift**0.5
  
      return schedule
  
imports: import collections
  import os
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import linen
  from flax.training import checkpoints
  from flax.training import common_utils
  import jax
  from jax import random
  import jax.numpy
  from jax.sharding import PartitionSpec, Mesh, NamedSharding
  import ml_collections
  import numpy
  import optax
  import tensorflow
  import input_pipeline
  import models
  import temperature_sampler
  import utils
label_desc: Use image node_id 1 for calling a global function with example usage: rsqrt_schedule(init_value, shift) and returns: schedule, unknown