code: def make_schedule(scalar_or_schedule: float | Schedule) -> Schedule:
      if callable(scalar_or_schedule):
          return scalar_or_schedule
      elif jnp.ndim(scalar_or_schedule) == 0:
          return constant(scalar_or_schedule)
      else:
          raise TypeError(type(scalar_or_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 16 for calling a global function with example usage: make_schedule(scalar_or_schedule) and returns: scalar_or_schedule, constant