code: def inverse_time_decay(
      step_size, decay_steps, decay_rate, staircase=False
  ):
      if staircase:
  
          def schedule(i):
              return step_size / (
                  1 + decay_rate * jnp.floor(i / decay_steps)
              )
  
      else:
  
          def schedule(i):
              return step_size / (1 + decay_rate * i / decay_steps)
  
      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 13 for calling a global function with example usage: inverse_time_decay(step_size, decay_steps, decay_rate, staircase) and returns: schedule, unknown, unknown