code: def beta_schedule(
      schedule, global_step, final_beta, decay_start, decay_end
  ):
      """Get KL multiplier (beta) based on the schedule."""
      if decay_start > decay_end:
          raise ValueError("decay_end is smaller than decay_end.")
  
      # Since some of the TF schedules do not support incrementing a value,
      # in all of the schedules, we anneal the beta from final_beta to zero
      # and then reverse it at the bottom.
      if schedule == "constant":
          decayed_value = 0.0
      elif schedule == "linear":
          decayed_value = tf.train.polynomial_decay(
              learning_rate=final_beta,
              global_step=global_step - decay_start,
              decay_steps=decay_end - decay_start,
              end_learning_rate=0.0,
          )
      elif schedule == "noisy_linear_cosine_decay":
          decayed_value = tf.train.noisy_linear_cosine_decay(
              learning_rate=final_beta,
              global_step=global_step - decay_start,
              decay_steps=decay_end - decay_start,
          )
      # TODO(mechcoder): Add log_annealing schedule.
      else:
          raise ValueError("Unknown beta schedule.")
  
      increased_value = final_beta - decayed_value
      increased_value = tf.maximum(0.0, increased_value)
  
      beta = tf.case(
          pred_fn_pairs={
              tf.less(global_step, decay_start): lambda: 0.0,
              tf.greater(global_step, decay_end): lambda: final_beta,
          },
          default=lambda: increased_value,
      )
      return beta
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import numpy
  from tensor2tensor.layers import common_layers
  from tensor2tensor.utils import contrib
  import tensorflow.compat.v1
  from tensorflow.python.ops import summary_op_util
label_desc: Use image node_id 21 for calling a global function with example usage: beta_schedule(schedule, global_step, final_beta, decay_start, decay_end) and returns: beta