code: def _polynomial_decay(
      step: fun_mc.AnyArray,
      step_size: fun_mc.FloatArray,
      decay_steps: fun_mc.AnyArray,
      final_step_size: fun_mc.FloatArray,
      power: float | fun_mc.FloatArray = 1.0,
  ) -> fun_mc.FloatArray:
      """Polynomial decay step size schedule."""
      step_size = jnp.asarray(step_size)
      step_f = jnp.array(step, step_size.dtype)
      decay_steps_f = jnp.array(decay_steps, step_size.dtype)
      step_mult = (1.0 - step_f / decay_steps_f) ** power
      step_mult = jnp.where(
          step >= decay_steps, jnp.zeros_like(step_mult), step_mult
      )
      return step_mult * (step_size - final_step_size) + final_step_size
  
imports: import functools
  import typing
  from typing import Any, Callable, Iterable, NamedTuple, Optional
  import numpy
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import malt
  from fun_mc import sga_hmc
  from fun_mc.malt import *
  from fun_mc.sga_hmc import *
label_desc: Use image node_id 1 for calling a global function with example usage: _polynomial_decay(step, step_size, decay_steps, final_step_size, power) and returns: unknown