code: def create_learning_rate_schedule(
      learning_rate: float, warmup_steps: int
  ):
      """Creates a rsqrt schedule with linear warmup."""
      return optax.join_schedules(
          [
              optax.linear_schedule(
                  init_value=0,
                  end_value=learning_rate,
                  transition_steps=warmup_steps,
              ),
              rsqrt_schedule(
                  init_value=learning_rate, shift=warmup_steps
              ),
          ],
          boundaries=[warmup_steps],
      )
  
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 2 for calling a global function with example usage: create_learning_rate_schedule(learning_rate, warmup_steps) and returns: optax