code: def _swap_discharge(x, val, idx, tree):
      indexers = tree_util.tree_unflatten(tree, idx)
      if len(indexers) > 1:
          raise NotImplementedError("Only single indexer is supported.")
      indexer = indexers[0]
      if _is_trivial_indexer(indexer):
          return x, val
      # If everything in the indexer is a slice or ()-shaped, we can also
      # use `lax.dynamic_slice` with 1-sized slices for ()-shaped indices.
      # We need to squeeze out the the 1-sized slices at the end.
      if maybe_slice := _maybe_convert_to_dynamic_slice(indexer):
          starts, sizes, squeeze_dims = maybe_slice
          x_old = lax_slicing.dynamic_slice(x, starts, sizes)
          val = lax.expand_dims(val, squeeze_dims)
          y = lax_slicing.dynamic_update_slice(x, val, starts)
          return lax.squeeze(x_old, squeeze_dims), y
      indexer = _convert_to_array_indexer(indexer)
      x_old = _prepend_gather(x, indexer)
      return x_old, _prepend_scatter(x, indexer, val)
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import dataclasses
  from functools import partial
  import operator
  from typing import Any, Callable, Protocol
  import numpy
  from jax._src import api_util
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import tree_util
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.lax import lax
  from jax._src.lax import slicing
  from jax._src.state import indexing
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.state.primitives import get_p, swap_p, addupdate_p
  from jax._src.state.utils import hoist_consts_to_refs
  from jax._src.typing import Array
  from jax._src.util import safe_map, safe_zip, split_list, weakref_lru_cache, partition_list, merge_lists, split_dict
label_desc: Use image node_id 14 for calling a global function with example usage: _swap_discharge(x, val, idx, tree) and returns: x_old,  _prepend_scatter, x,  val, lax,  y