code: def create_learning_rate_fn(
      train_ds_size: int,
      train_batch_size: int,
      num_train_epochs: int,
      num_warmup_steps: int,
      learning_rate: float,
  ) -> Callable[[int], jnp.ndarray]:
      """Returns a linear warmup, linear_decay learning rate function."""
      steps_per_epoch = train_ds_size // train_batch_size
      num_train_steps = steps_per_epoch * num_train_epochs
      warmup_fn = optax.linear_schedule(
          init_value=0.0,
          end_value=learning_rate,
          transition_steps=num_warmup_steps,
      )
      decay_fn = optax.linear_schedule(
          init_value=learning_rate,
          end_value=0,
          transition_steps=num_train_steps - num_warmup_steps,
      )
      schedule_fn = optax.join_schedules(
          schedules=[warmup_fn, decay_fn], boundaries=[num_warmup_steps]
      )
      return schedule_fn
  
imports: import logging
  import math
  import os
  import sys
  import time
  from dataclasses import dataclass, field
  from itertools import chain
  from pathlib import Path
  from typing import Callable, Optional
  import datasets
  import jax
  import jax.numpy
  import numpy
  import optax
  from datasets import Dataset, load_dataset
  from flax.core.frozen_dict import freeze, unfreeze
  from flax.training.common_utils import onehot, stack_forest
  from jax.experimental.maps import mesh
  from jax.experimental.pjit import pjit
  from partitions import set_partitions
  from tqdm import tqdm
  import transformers
  from transformers import CONFIG_MAPPING, FLAX_MODEL_FOR_CAUSAL_LM_MAPPING, AutoConfig, AutoTokenizer, FlaxAutoModelForCausalLM, HfArgumentParser, TrainingArguments, is_tensorboard_available
  from transformers.testing_utils import CaptureLogger
label_desc: Use image node_id 4 for calling a global function with example usage: create_learning_rate_fn(train_ds_size, train_batch_size, num_train_epochs, num_warmup_steps, learning_rate) and returns: schedule_fn