code: def semaphore_wait(sem_or_view, dec: int | jax.Array = 1):
      ref, indexers = _get_ref_and_indexers(sem_or_view)
      dec = jnp.asarray(dec, dtype=jnp.int32)
      args = [ref, indexers, dec]
      flat_args, args_tree = tree_util.tree_flatten(args)
      semaphore_wait_p.bind(*flat_args, args_tree=args_tree)
  
imports: from __future__ import annotations
  import contextlib
  import dataclasses
  import enum
  from typing import Any, Callable
  import jax
  from jax._src import api_util
  from jax._src import core
  from jax._src import effects
  from jax._src import linear_util
  from jax._src import pretty_printer
  from jax._src import state
  from jax._src import tree_util
  from jax._src import util
  from jax._src.state import indexing
  from jax._src.state import primitives
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.pallas.mosaic import core
  import jax.numpy
label_desc: Use image node_id 14 for calling a global function with example usage: semaphore_wait(sem_or_view, dec) without return types