code: def piecewise_constant(boundaries: Any, values: Any):
      boundaries = jnp.array(boundaries)
      values = jnp.array(values)
      if not boundaries.ndim == values.ndim == 1:
          raise ValueError("boundaries and values must be sequences")
      if not boundaries.shape[0] == values.shape[0] - 1:
          raise ValueError(
              "boundaries length must be one shorter than values length"
          )
  
      def schedule(i):
          return values[jnp.sum(i > boundaries)]
  
      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 15 for calling a global function with example usage: piecewise_constant(boundaries, values) and returns: schedule, values