code: def _get_callback(treedef, static, is_float0):
      @ft.partial(jax.jit, static_argnums=0)
      def callback(static_fn, dynamic):
          # pure_callback casts float0 to bool, so need to cast back
          assert len(dynamic) == len(is_float0)
          dynamic = [
              np.broadcast_to(
                  np.zeros((), dtype=jax.dtypes.float0), x.shape
              )
              if y
              else x
              for x, y in zip(dynamic, is_float0)
          ]
          iter_dynamic = iter(dynamic)
          flat = [
              next(iter_dynamic) if x is None else x for x in static
          ]
          assert next(iter_dynamic, None) is None
          transforms, args = jtu.tree_unflatten(treedef, flat)
          for transform in transforms:
              static_fn = transform(static_fn)
          out = static_fn(args)
          return tuple(jtu.tree_leaves(out))
  
      def callback_lookup(dynamic_index, *dynamic):
          static_fn = _index_to_fn[dynamic_index.item()]
          return callback(static_fn, dynamic)
  
      return callback_lookup
  
imports: import functools
  from collections.abc import Callable
  from typing import Any, Optional, Union
  import jax
  import jax._src.pretty_printer
  import jax._src.source_info_util
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.numpy
  import jax.tree_util
  import numpy
  from jaxtyping import Array, Int, PyTree
  from .._ad import filter_jvp
  from .._caches import internal_lru_caches
  from .._compile_utils import hashable_combine, hashable_partition
  from .._eval_shape import filter_eval_shape
  from .._filters import combine, is_array, partition
  from .._module import field, Module, module_update_wrapper
  from .._vmap_pmap import filter_vmap
  from .None import _primitive
  from ._primitive import filter_primitive_batching, filter_primitive_bind, filter_primitive_def, filter_primitive_jvp, filter_primitive_transpose, materialise_zeros
label_desc: Use image node_id 10 for calling a global function with example usage: _get_callback(treedef, static, is_float0) and returns: callback_lookup, tuple, callback