code: def create_learning_rate_fn(
      config: ml_collections.ConfigDict,
      base_learning_rate: float,
      steps_per_epoch: int,
  ):
      """Create learning rate schedule."""
      warmup_fn = optax.linear_schedule(
          init_value=0.0,
          end_value=base_learning_rate,
          transition_steps=config.warmup_epochs * steps_per_epoch,
      )
      cosine_epochs = max(config.num_epochs - config.warmup_epochs, 1)
      cosine_fn = optax.cosine_decay_schedule(
          init_value=base_learning_rate,
          decay_steps=cosine_epochs * steps_per_epoch,
      )
      schedule_fn = optax.join_schedules(
          schedules=[warmup_fn, cosine_fn],
          boundaries=[config.warmup_epochs * steps_per_epoch],
      )
      return schedule_fn
  
imports: import functools
  import time
  from typing import Any
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax.training import checkpoints
  from flax.training import common_utils
  from flax.training import dynamic_scale
  from flax.training import train_state
  import jax
  from jax import lax
  import jax.numpy
  from jax import random
  import ml_collections
  import optax
  import tensorflow
  import tensorflow_datasets
  import input_pipeline
  import models
label_desc: Use image node_id 5 for calling a global function with example usage: create_learning_rate_fn(config, base_learning_rate, steps_per_epoch) and returns: schedule_fn