code: def _scatter_update(
      x,
      idx,
      y,
      scatter_op,
      indices_are_sorted,
      unique_indices,
      mode=None,
      normalize_indices=True,
  ):
      """Helper for indexed updates.
  
      Computes the value of x that would result from computing::
        x[idx] op= y
      except in a pure functional way, with no in-place updating.
  
      Args:
        x: ndarray to be updated.
        idx: None, an integer, a slice, an ellipsis, an ndarray with integer dtype,
          or a tuple of those indicating the locations of `x` into which to scatter-
          update the values in `y`.
        y: values to be scattered.
        scatter_op: callable, one of lax.scatter, lax.scatter_add, lax.scatter_min,
          or lax_scatter_max.
        indices_are_sorted: whether `idx` is known to be sorted
        unique_indices: whether `idx` is known to be free of duplicates
  
      Returns:
        An ndarray representing an updated `x` after performing the scatter-update.
      """
      x = jnp.asarray(x)
      if (
          isinstance(y, int)
          and np.issubdtype(x.dtype, np.integer)
          and np.iinfo(x.dtype).min <= y <= np.iinfo(x.dtype).max
      ):
          y = jnp.asarray(y, dtype=x.dtype)
      else:
          y = jnp.asarray(y)
  
      # XLA gathers and scatters are very similar in structure; the scatter logic
      # is more or less a transpose of the gather equivalent.
      treedef, static_idx, dynamic_idx = jnp._split_index_for_jit(
          idx, x.shape
      )
      return _scatter_impl(
          x,
          y,
          scatter_op,
          treedef,
          static_idx,
          dynamic_idx,
          indices_are_sorted,
          unique_indices,
          mode,
          normalize_indices,
      )
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import sys
  from typing import Callable, Union
  import warnings
  import numpy
  from jax import lax
  from jax._src import config
  from jax._src import core
  from jax._src import dtypes
  from jax._src import util
  from jax._src.lax import lax
  from jax._src.numpy import lax_numpy
  from jax._src.numpy import reductions
  from jax._src.numpy.util import check_arraylike, promote_dtypes
  from jax._src.typing import Array, ArrayLike
label_desc: Use image node_id 1 for calling a global function with example usage: _scatter_update(x, idx, y, scatter_op, indices_are_sorted, unique_indices, mode, normalize_indices) and returns: _scatter_impl