code: def gpt3_schedule(warmup_steps, total_steps, peak_lr, end_lr):
      def sch(step):
          warmup_pct = jnp.clip(step, 0, warmup_steps) / warmup_steps
          anneal_pct = (
              jnp.clip(step - warmup_steps, 0, total_steps)
              / total_steps
          )
  
          return (
              warmup_pct * peak_lr
              - (peak_lr - end_lr)
              * (1 - jnp.cos(jnp.pi * anneal_pct))
              / 2
          )
  
      return sch
  
imports: import jax
  import jax.numpy
  from jax.experimental.pjit import with_sharding_constraint
  from optax import AdditiveWeightDecayState, GradientTransformation, OptState
label_desc: Use image node_id 2 for calling a global function with example usage: gpt3_schedule(warmup_steps, total_steps, peak_lr, end_lr) and returns: sch, unknown