code: def debug_print(
      fmt: str, *args, ordered: bool = False, **kwargs
  ) -> None:
      """Prints values and works in staged out JAX functions.
  
      This function does *not* work with f-strings because formatting is delayed.
      So instead of ``jax.debug.print(f"hello {bar}")``, write
      ``jax.debug.print("hello {bar}", bar=bar)``.
  
      This function is a thin convenience wrapper around :func:`jax.debug.callback`.
      The implementation is essentially::
  
        def debug_print(fmt: str, *args, **kwargs):
          jax.debug.callback(
              lambda *args, **kwargs: print(fmt.format(*args, **kwargs)),
              *args, **kwargs)
  
      It may be useful to call :func:`jax.debug.callback` directly instead of this
      convenience wrapper. For example, to get debug printing in logs, you might
      use :func:`jax.debug.callback` together with ``logging.log``.
  
      Args:
        fmt: A format string, e.g. ``"hello {x}"``, that will be used to format
          input arguments, like ``str.format``. See the Python docs on
          `string formatting <https://docs.python.org/3/library/stdtypes.html#str.format>`_
          and `format string syntax <https://docs.python.org/3/library/string.html#formatstrings>`_.
        *args: A list of positional arguments to be formatted, as if passed to
          ``fmt.format``.
        ordered: A keyword only argument used to indicate whether or not the
          staged out computation will enforce ordering of this ``jax.debug.print``
          w.r.t. other ordered ``jax.debug.print`` calls.
        **kwargs: Additional keyword arguments to be formatted, as if passed to
          ``fmt.format``.
      """
      # Check that we provide the correct arguments to be formatted
      formatter.format(fmt, *args, **kwargs)
  
      debug_callback(
          functools.partial(_format_print_callback, fmt),
          *args,
          **kwargs,
          ordered=ordered,
      )
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import string
  import sys
  from typing import Any, Callable, Union
  import weakref
  import numpy
  import jax.numpy
  from jax import lax
  from jax._src import core
  from jax._src import effects
  from jax._src import mesh
  from jax._src import sharding_impls
  from jax._src import dispatch
  from jax._src import tree_util
  from jax._src import util
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.lib import xla_client
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  from jax._src.sharding import Sharding
  from jax._src.sharding_impls import NamedSharding, parse_flatten_op_sharding
label_desc: Use image node_id 10 for calling a global function with example usage: debug_print(fmt) without return types