code: def polynomial_decay(
      step_size, decay_steps, final_step_size, power=1.0
  ):
      def schedule(step_num):
          step_num = jnp.minimum(step_num, decay_steps)
          step_mult = (1 - step_num / decay_steps) ** power
          return (
              step_mult * (step_size - final_step_size)
              + final_step_size
          )
  
      return schedule
  
imports: from __future__ import annotations
  from typing import Any, Callable, NamedTuple
  from collections import namedtuple
  import functools
  from functools import partial
  import jax.numpy
  from jax._src.util import safe_zip, safe_map, unzip2
  from jax import tree_util
  from jax.tree_util import tree_map, tree_flatten, tree_unflatten, register_pytree_node
label_desc: Use image node_id 14 for calling a global function with example usage: polynomial_decay(step_size, decay_steps, final_step_size, power) and returns: schedule, unknown