code: def pallas_call(
      f: Callable[..., None],
      out_shape: Any,
      *,
      grid_spec: GridSpec | None = None,
      debug: bool = False,
      grid: Grid | None = None,
      in_specs: Sequence[BlockSpec | NoBlockSpec]
      | NoBlockSpec = no_block_spec,
      out_specs: BlockSpec
      | NoBlockSpec
      | Sequence[BlockSpec | NoBlockSpec] = no_block_spec,
      input_output_aliases: dict[int, int] = {},
      interpret: bool = False,
      name: str | None = None,
      **compiler_params: Any,
  ):
      name = _extract_function_name(f, name)
      if grid_spec is None:
          grid_spec = GridSpec(grid, in_specs, out_specs)
      if isinstance(out_shape, list):
          out_shape = tuple(out_shape)
      flat_out_shapes, out_tree = tree_util.tree_flatten(out_shape)
      flat_out_shapes = [
          jax.ShapeDtypeStruct(x.shape, x.dtype)
          for x in flat_out_shapes
      ]
  
      @jax.jit
      def wrapped(*args):
          flat_args, in_tree = tree_util.tree_flatten(args)
          flat_in_avals = tuple(
              jax_core.raise_to_shaped(jax_core.get_aval(a))
              for a in flat_args
          )
          flat_out_avals = tuple(
              jax_core.ShapedArray(v.shape, v.dtype)
              for v in flat_out_shapes
          )
          grid_mapping, jaxpr, consts, _ = _trace_to_jaxpr(
              f,
              grid_spec,
              flat_in_avals,
              flat_out_avals,
              in_tree,
              out_tree,
          )
          which_linear = (False,) * len(flat_args)
          out_flat = pallas_call_p.bind(
              *consts,
              *flat_args,
              jaxpr=jaxpr,
              name=name,
              which_linear=which_linear,
              in_shapes=tuple(
                  jax.ShapeDtypeStruct(a.shape, a.dtype)
                  for a in flat_args
              ),
              out_shapes=tuple(flat_out_shapes),
              debug=debug,
              interpret=interpret,
              grid_mapping=grid_mapping,
              input_output_aliases=tuple(input_output_aliases.items()),
              **compiler_params,
          )
          out = tree_util.tree_unflatten(out_tree, out_flat)
          return out
  
      return wrapped
  
imports: from __future__ import annotations
  from functools import partial
  import itertools
  from typing import Any, Callable
  from collections.abc import Sequence
  import jax
  from jax import api_util
  from jax import tree_util
  from jax import lax
  from jax._src import state
  from jax.interpreters import ad
  from jax.interpreters import batching
  from jax.interpreters import partial_eval
  from jax.interpreters import mlir
  from jax.interpreters import xla
  from jax._src import ad_util
  from jax._src import core
  from jax._src.state import primitives
  from jax._src import linear_util
  from jax._src.state import discharge
  from jax._src.util import split_list, safe_map, safe_zip, weakref_lru_cache, tuple_insert, partition_list, merge_lists
  import jax.numpy
  import numpy
  from jax._src.pallas import core
label_desc: Use image node_id 13 for calling a global function with example usage: pallas_call(f, out_shape) and returns: wrapped, out