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.array]:
      """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 json
  import logging
  import math
  import os
  import random
  import sys
  import time
  from dataclasses import dataclass, field
  from pathlib import Path
  from typing import Any, Callable, Dict, Optional, Tuple
  import datasets
  import numpy
  from datasets import load_dataset
  from tqdm import tqdm
  import evaluate
  import jax
  import jax.numpy
  import optax
  import transformers
  from flax import struct, traverse_util
  from flax.jax_utils import pad_shard_unpad, replicate, unreplicate
  from flax.training import train_state
  from flax.training.common_utils import get_metrics, onehot, shard
  from huggingface_hub import Repository
  from transformers import AutoConfig, AutoTokenizer, FlaxAutoModelForSequenceClassification, HfArgumentParser, PretrainedConfig, TrainingArguments, is_tensorboard_available
  from transformers.utils import check_min_version, get_full_repo_name, send_example_telemetry
label_desc: Use image node_id 2 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