# Copyright 2018 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from collections.abc import Callable, Sequence
import enum
import operator
from functools import partial
import math
from typing import NamedTuple
import weakref

import numpy as np

from jax._src import ad_util
from jax._src import api
from jax._src import core
from jax._src import dispatch
from jax._src import dtypes
from jax._src import source_info_util
from jax._src.traceback_util import api_boundary
from jax._src import util
from jax._src import mesh as mesh_lib
from jax._src.interpreters import ad
from jax._src.interpreters import batching
from jax._src.interpreters import mlir
from jax._src.interpreters import partial_eval as pe
from jax._src.lax import lax
from jax._src.lax import utils as lax_utils
from jax._src.lax.utils import (
    _argnum_weak_type,
    input_dtype,
    standard_primitive,
)
from jax._src.lib.mlir import ir
from jax._src.lib.mlir.dialects import hlo
from jax._src.named_sharding import NamedSharding
from jax._src.partition_spec import PartitionSpec as P
from jax._src.typing import Array, ArrayLike, Shape
from jax._src.state.indexing import ds
from jax._src.util import safe_map, safe_zip

map, unsafe_map = safe_map, map
zip, unsafe_zip = safe_zip, zip

_dtype = dtypes.dtype
_slice = slice


def slice(operand: ArrayLike, start_indices: Sequence[int],
          limit_indices: Sequence[int],
          strides: Sequence[int] | None = None) -> Array:
  """Wraps XLA's `Slice
  <https://www.openxla.org/xla/operation_semantics#slice>`_
  operator.

  Args:
    operand: an array to slice
    start_indices: a sequence of ``operand.ndim`` start indices.
    limit_indices: a sequence of ``operand.ndim`` limit indices.
    strides: an optional sequence of ``operand.ndim`` strides.

  Returns:
    The sliced array

  Examples:
    Here are some examples of simple two-dimensional slices:

    >>> x = jnp.arange(12).reshape(3, 4)
    >>> x
    Array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]], dtype=int32)

    >>> lax.slice(x, (1, 0), (3, 2))
    Array([[4, 5],
           [8, 9]], dtype=int32)

    >>> lax.slice(x, (0, 0), (3, 4), (1, 2))
    Array([[ 0,  2],
           [ 4,  6],
           [ 8, 10]], dtype=int32)

    These two examples are equivalent to the following Python slicing syntax:

    >>> x[1:3, 0:2]
    Array([[4, 5],
           [8, 9]], dtype=int32)

    >>> x[0:3, 0:4:2]
    Array([[ 0,  2],
           [ 4,  6],
           [ 8, 10]], dtype=int32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :func:`jax.lax.slice_in_dim`
    - :func:`jax.lax.index_in_dim`
    - :func:`jax.lax.dynamic_slice`
  """
  return slice_p.bind(operand, start_indices=tuple(start_indices),
                      limit_indices=tuple(limit_indices),
                      strides=None if strides is None else tuple(strides))


def dynamic_slice(
    operand: Array | np.ndarray,
    start_indices: Array | np.ndarray | Sequence[ArrayLike],
    slice_sizes: Shape,
    *,
    allow_negative_indices: bool | Sequence[bool] = True
) -> Array:
  """Wraps XLA's `DynamicSlice
  <https://www.openxla.org/xla/operation_semantics#dynamicslice>`_
  operator.

  Args:
    operand: an array to slice.
    start_indices: a list of scalar indices, one per dimension. These values
      may be dynamic.
    slice_sizes: the size of the slice. Must be a sequence of non-negative
      integers with length equal to `ndim(operand)`. Inside a JIT compiled
      function, only static values are supported (all JAX arrays inside JIT
      must have statically known size).
    allow_negative_indices: a bool or sequence of bools, one per dimension; if
      a bool is passed, it applies to all dimensions. For each dimension,
      if true, negative indices are permitted and are are interpreted relative
      to the end of the array. If false, negative indices are treated as if they
      were out of bounds and the result is implementation defined, typically
      clamped to the first index.

  Returns:
    An array containing the slice.

  Examples:
    Here is a simple two-dimensional dynamic slice:

    >>> x = jnp.arange(12).reshape(3, 4)
    >>> x
    Array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]], dtype=int32)

    >>> dynamic_slice(x, (1, 1), (2, 3))
    Array([[ 5,  6,  7],
           [ 9, 10, 11]], dtype=int32)

    Note the potentially surprising behavior for the case where the requested slice
    overruns the bounds of the array; in this case the start index is adjusted to
    return a slice of the requested size:

    >>> dynamic_slice(x, (1, 1), (2, 4))
    Array([[ 4,  5,  6,  7],
           [ 8,  9, 10, 11]], dtype=int32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :func:`jax.lax.slice`
    - :func:`jax.lax.dynamic_slice_in_dim`
    - :func:`jax.lax.dynamic_index_in_dim`
  """
  start_indices = _dynamic_slice_indices(
      operand, start_indices, allow_negative_indices)
  sizes = core.canonicalize_shape(slice_sizes)  # type: ignore
  operand, *start_indices = core.standard_insert_pvary(operand, *start_indices)
  return dynamic_slice_p.bind(operand, *start_indices, slice_sizes=tuple(sizes))


def dynamic_update_slice(
    operand: Array | np.ndarray, update: ArrayLike,
    start_indices: Array | Sequence[ArrayLike],
    *,
    allow_negative_indices: bool | Sequence[bool] = True
) -> Array:
  """Wraps XLA's `DynamicUpdateSlice
  <https://www.openxla.org/xla/operation_semantics#dynamicupdateslice>`_
  operator.

  Args:
    operand: an array to slice.
    update: an array containing the new values to write onto `operand`.
    start_indices: a list of scalar indices, one per dimension.
    allow_negative_indices: a bool or sequence of bools, one per dimension; if
      a bool is passed, it applies to all dimensions. For each dimension,
      if true, negative indices are permitted and are are interpreted relative
      to the end of the array. If false, negative indices are treated as if they
      were out of bounds and the result is implementation defined, typically
      clamped to the first index.

  Returns:
    An array containing the slice.

  Examples:
    Here is an example of updating a one-dimensional slice update:

    >>> x = jnp.zeros(6)
    >>> y = jnp.ones(3)
    >>> dynamic_update_slice(x, y, (2,))
    Array([0., 0., 1., 1., 1., 0.], dtype=float32)

    If the update slice is too large to fit in the array, the start
    index will be adjusted to make it fit

    >>> dynamic_update_slice(x, y, (3,))
    Array([0., 0., 0., 1., 1., 1.], dtype=float32)
    >>> dynamic_update_slice(x, y, (5,))
    Array([0., 0., 0., 1., 1., 1.], dtype=float32)

    Here is an example of a two-dimensional slice update:

    >>> x = jnp.zeros((4, 4))
    >>> y = jnp.ones((2, 2))
    >>> dynamic_update_slice(x, y, (1, 2))
    Array([[0., 0., 0., 0.],
           [0., 0., 1., 1.],
           [0., 0., 1., 1.],
           [0., 0., 0., 0.]], dtype=float32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :attr:`lax.dynamic_update_index_in_dim`
    - :attr:`lax.dynamic_update_slice_in_dim`
  """
  start_indices = _dynamic_slice_indices(
      operand, start_indices, allow_negative_indices)
  operand, update, *start_indices = core.standard_insert_pvary(
      operand, update, *start_indices)
  return dynamic_update_slice_p.bind(operand, update, *start_indices)


class GatherDimensionNumbers(NamedTuple):
  """
  Describes the dimension number arguments to an `XLA's Gather operator
  <https://www.openxla.org/xla/operation_semantics#gather>`_. See the XLA
  documentation for more details of what the dimension numbers mean.

  Args:
    offset_dims: the set of dimensions in the `gather` output that offset into
      an array sliced from `operand`. Must be a tuple of integers in ascending
      order, each representing a dimension number of the output.
    collapsed_slice_dims: the set of dimensions `i` in `operand` that have
      `slice_sizes[i] == 1` and that should not have a corresponding dimension
      in the output of the gather. Must be a tuple of integers in ascending
      order.
    start_index_map: for each dimension in `start_indices`, gives the
      corresponding dimension in the `operand` that is to be sliced. Must be a
      tuple of integers with size equal to `start_indices.shape[-1]`.
    operand_batching_dims: the set of batching dimensions `i` in `operand` that
      have `slice_sizes[i] == 1` and that should have a corresponding dimension
      in both the `start_indices` (at the same index in
      `start_indices_batching_dims`) and output of the gather. Must be a tuple
      of integers in ascending order.
    start_indices_batching_dims: the set of batching dimensions `i` in
      `start_indices` that should have a corresponding dimension in both the
      `operand` (at the same index in `operand_batching_dims`) and output of the
      gather. Must be a tuple of integers (order is fixed based on
      correspondence with `operand_batching_dims`).

  Unlike XLA's `GatherDimensionNumbers` structure, `index_vector_dim` is
  implicit; there is always an index vector dimension and it must always be the
  last dimension. To gather scalar indices, add a trailing dimension of size 1.
  """
  offset_dims: tuple[int, ...]
  collapsed_slice_dims: tuple[int, ...]
  start_index_map: tuple[int, ...]
  operand_batching_dims: tuple[int, ...] = ()
  start_indices_batching_dims: tuple[int, ...] = ()


class GatherScatterMode(enum.Enum):
  """
  Describes how to handle out-of-bounds indices in a gather or scatter.

  Possible values are:

  CLIP:
    Indices will be clamped to the nearest in-range value, i.e., such that the
    entire window to be gathered is in-range.
  FILL_OR_DROP:
    If any part of a gathered window is out of bounds, the entire window
    that is returned, even those elements that were otherwise in-bounds, will be
    filled with a constant.
    If any part of a scattered window is out of bounds, the entire window
    will be discarded.
  PROMISE_IN_BOUNDS:
    The user promises that indices are in bounds. No additional checking will be
    performed. In practice, with the current XLA implementation this means
    that out-of-bounds gathers will be clamped but out-of-bounds scatters will
    be discarded. Gradients will not be correct if indices are out-of-bounds.
  """
  CLIP = enum.auto()
  FILL_OR_DROP = enum.auto()
  PROMISE_IN_BOUNDS = enum.auto()
  ONE_HOT = enum.auto()

  @staticmethod
  def from_any(s: str | GatherScatterMode | None) -> GatherScatterMode:
    if isinstance(s, GatherScatterMode):
      return s
    if s == "clip":
      return GatherScatterMode.CLIP
    if s is None or s == "fill" or s == "drop":
      return GatherScatterMode.FILL_OR_DROP
    if s == "promise_in_bounds":
      return GatherScatterMode.PROMISE_IN_BOUNDS
    if s == "one_hot":
      return GatherScatterMode.ONE_HOT
    else:
      raise ValueError(f'Unknown gather mode "{s}"')


def gather(operand: ArrayLike, start_indices: ArrayLike,
           dimension_numbers: GatherDimensionNumbers,
           slice_sizes: Shape,
           *,
           unique_indices: bool = False,
           indices_are_sorted: bool = False,
           mode: str | GatherScatterMode | None = None,
           fill_value = None) -> Array:
  """Gather operator.

  Wraps `XLA's Gather operator
  <https://www.openxla.org/xla/operation_semantics#gather>`_.

  :func:`gather` is a low-level operator with complicated semantics, and most JAX
  users will never need to call it directly. Instead, you should prefer using
  `Numpy-style indexing`_, and/or :func:`jax.numpy.ndarray.at`, perhaps in combination
  with :func:`jax.vmap`.

  Args:
    operand: an array from which slices should be taken
    start_indices: the indices at which slices should be taken
    dimension_numbers: a `lax.GatherDimensionNumbers` object that describes
      how dimensions of `operand`, `start_indices` and the output relate.
    slice_sizes: the size of each slice. Must be a sequence of non-negative
      integers with length equal to `ndim(operand)`.
    indices_are_sorted: whether `indices` is known to be sorted. If
      true, may improve performance on some backends.
    unique_indices: whether the elements gathered from ``operand`` are
      guaranteed not to overlap with each other. If ``True``, this may improve
      performance on some backends. JAX does not check this promise: if
      the elements overlap the behavior is undefined.
    mode: how to handle indices that are out of bounds: when set to ``'clip'``,
      indices are clamped so that the slice is within bounds, and when
      set to ``'fill'`` or ``'drop'`` gather returns a slice full of
      ``fill_value`` for the affected slice. The behavior for out-of-bounds
      indices when set to ``'promise_in_bounds'`` is implementation-defined.
    fill_value: the fill value to return for out-of-bounds slices when `mode`
      is ``'fill'``. Ignored otherwise. Defaults to ``NaN`` for inexact types,
      the largest negative value for signed types, the largest positive value
      for unsigned types, and ``True`` for booleans.

  Returns:
    An array containing the gather output.

  Examples:
    As mentioned above, you should basically never use :func:`gather` directly,
    and instead use NumPy-style indexing expressions to gather values from
    arrays.

    For example, here is how you can extract values at particular indices using
    straightforward indexing semantics, which will lower to XLA's Gather operator:

    >>> import jax.numpy as jnp
    >>> x = jnp.array([10, 11, 12])
    >>> indices = jnp.array([0, 1, 1, 2, 2, 2])

    >>> x[indices]
    Array([10, 11, 11, 12, 12, 12], dtype=int32)

    For control over settings like ``indices_are_sorted``, ``unique_indices``, ``mode``,
    and ``fill_value``, you can use the :attr:`jax.numpy.ndarray.at` syntax:

    >>> x.at[indices].get(indices_are_sorted=True, mode="promise_in_bounds")
    Array([10, 11, 11, 12, 12, 12], dtype=int32)

    By comparison, here is the equivalent function call using :func:`gather` directly,
    which is not something typical users should ever need to do:

    >>> from jax import lax
    >>> lax.gather(x, indices[:, None], slice_sizes=(1,),
    ...            dimension_numbers=lax.GatherDimensionNumbers(
    ...                offset_dims=(),
    ...                collapsed_slice_dims=(0,),
    ...                start_index_map=(0,)),
    ...            indices_are_sorted=True,
    ...            mode=lax.GatherScatterMode.PROMISE_IN_BOUNDS)
    Array([10, 11, 11, 12, 12, 12], dtype=int32)

  .. _Numpy-style indexing: https://numpy.org/doc/stable/reference/arrays.indexing.html
  """
  if mode is None:
    mode = GatherScatterMode.PROMISE_IN_BOUNDS
  parsed_mode = GatherScatterMode.from_any(mode)
  if parsed_mode == GatherScatterMode.FILL_OR_DROP:
    if fill_value is None:
      dtype = _dtype(operand)
      if dtypes.issubdtype(dtype, np.inexact):
        fill_value = np.nan
      elif dtypes.issubdtype(dtype, np.signedinteger):
        fill_value = dtypes.iinfo(dtype).min
      elif dtypes.issubdtype(dtype, np.unsignedinteger):
        fill_value = dtypes.iinfo(dtype).max
      elif dtype == dtypes.bool_:
        fill_value = True
      elif dtypes.issubdtype(dtype, dtypes.prng_key):
        fill_value = np.iinfo('uint32').max
      else:
        raise ValueError(f"Unsupported dtype for gather fill_value {dtype}")
  else:
    fill_value = None
  operand, start_indices = core.standard_insert_pvary(operand, start_indices)
  return gather_p.bind(
      operand, start_indices, dimension_numbers=dimension_numbers,
      slice_sizes=core.canonicalize_shape(slice_sizes),
      unique_indices=bool(unique_indices),
      indices_are_sorted=bool(indices_are_sorted),
      mode=parsed_mode,
      fill_value=fill_value)


class ScatterDimensionNumbers(NamedTuple):
  """
  Describes the dimension number arguments to an `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_. See the XLA
  documentation for more details of what the dimension numbers mean.

  Args:
    update_window_dims: the set of dimensions in the `updates` that are window
      dimensions. Must be a tuple of integers in ascending
      order, each representing a dimension number.
    inserted_window_dims: the set of size 1 window dimensions that must be
      inserted into the shape of `updates`. Must be a tuple of integers in
      ascending order, each representing a dimension number of the output. These
      are the mirror image of `collapsed_slice_dims` in the case of `gather`.
    scatter_dims_to_operand_dims: for each dimension in `scatter_indices`, gives
      the corresponding dimension in `operand`. Must be a sequence of integers
      with size equal to `scatter_indices.shape[-1]`.
    operand_batching_dims: the set of batching dimensions `i` in `operand` that
      should have a corresponding dimension in both the `scatter_indices` (at
      the same index in `scatter_indices_batching_dims`) and `updates`. Must be
      a tuple of integers in ascending order. These are the mirror image of
      `operand_batching_dims` in the case of `gather`.
    scatter_indices_batching_dims: the set of batching dimensions `i` in
      `scatter_indices` that should have a corresponding dimension in both the
      `operand` (at the same index in `operand_batching_dims`) and output of the
      gather. Must be a tuple of integers (order is fixed based on
      correspondence with `input_batching_dims`). These are the mirror image of
      `start_indices_batching_dims` in the case of `gather`.

  Unlike XLA's `ScatterDimensionNumbers` structure, `index_vector_dim` is
  implicit; there is always an index vector dimension and it must always be the
  last dimension. To scatter scalar indices, add a trailing dimension of size 1.
  """
  update_window_dims: Sequence[int]
  inserted_window_dims: Sequence[int]
  scatter_dims_to_operand_dims: Sequence[int]
  operand_batching_dims: Sequence[int] = ()
  scatter_indices_batching_dims: Sequence[int] = ()

@partial(api_boundary, repro_api_name="lax.scatter_add")
def scatter_add(
  operand: ArrayLike, scatter_indices: ArrayLike, updates: ArrayLike,
  dimension_numbers: ScatterDimensionNumbers, *,
  indices_are_sorted: bool = False, unique_indices: bool = False,
  mode: str | GatherScatterMode | None = None) -> Array:
  """Scatter-add operator.

  Wraps `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_, where
  addition is used to combine updates and values from `operand`.

  The semantics of scatter are complicated, and its API might change in the
  future. For most use cases, you should prefer the
  :attr:`jax.numpy.ndarray.at` property on JAX arrays which uses
  the familiar NumPy indexing syntax.

  Args:
    operand: an array to which the scatter should be applied
    scatter_indices: an array that gives the indices in `operand` to which each
      update in `updates` should be applied.
    updates: the updates that should be scattered onto `operand`.
    dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes how
      dimensions of `operand`, `scatter_indices`, `updates` and the output
      relate.
    indices_are_sorted: whether `scatter_indices` is known to be sorted. If
      true, may improve performance on some backends.
    unique_indices: whether the elements to be updated in ``operand`` are
      guaranteed to not overlap with each other. If true, may improve
      performance on some backends. JAX does not check this promise: if the
      updated elements overlap when ``unique_indices`` is ``True`` the behavior
      is undefined.
    mode: how to handle indices that are out of bounds: when set to 'clip',
      indices are clamped so that the slice is within bounds, and when set to
      'fill' or 'drop' out-of-bounds updates are dropped. The behavior for
      out-of-bounds indices when set to 'promise_in_bounds' is
      implementation-defined.

  Returns:
    An array containing the sum of `operand` and the scattered updates.

  Examples:
    As mentioned above, you should basically never use :func:`scatter_add`
    directly, and instead perform scatter-style operations using NumPy-style
    indexing expressions via :attr:`jax.numpy.ndarray.at`.

    Here is and example of updating entries in an array using
    :attr:`jax.numpy.ndarray.at`, which lowers to an XLA Scatter operation:

    >>> x = jnp.ones(5)
    >>> indices = jnp.array([1, 2, 4])
    >>> values = jnp.array([2.0, 3.0, 4.0])

    >>> x.at[indices].add(values)
    Array([1., 3., 4., 1., 5.], dtype=float32)

    This syntax also supports several of the optional arguments to
    :func:`scatter_add`, for example:

    >>> x.at[indices].add(values, indices_are_sorted=True,
    ...                   mode='promise_in_bounds')
    Array([1., 3., 4., 1., 5.], dtype=float32)

    By comparison, here is the equivalent function call using
    :func:`scatter_add` directly, which is not something typical users should
    ever need to do:

    >>> lax.scatter_add(x, indices[:, None], values,
    ...                 dimension_numbers=lax.ScatterDimensionNumbers(
    ...                     update_window_dims=(),
    ...                     inserted_window_dims=(0,),
    ...                     scatter_dims_to_operand_dims=(0,)),
    ...                 indices_are_sorted=True,
    ...                 mode=lax.GatherScatterMode.PROMISE_IN_BOUNDS)
    Array([1., 3., 4., 1., 5.], dtype=float32)
  """
  jaxpr, consts = lax._reduction_jaxpr(lax.add,
                                       core.get_aval(lax._const(operand, 0)))
  operand, scatter_indices, updates = core.standard_insert_pvary(
      operand, scatter_indices, updates)
  return scatter_add_p.bind(
      operand, scatter_indices, updates, update_jaxpr=jaxpr,
      update_consts=consts, dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=GatherScatterMode.from_any(mode))

@partial(api_boundary, repro_api_name="lax.scatter_sub")
def scatter_sub(
    operand: ArrayLike,
    scatter_indices: ArrayLike,
    updates: ArrayLike,
    dimension_numbers: ScatterDimensionNumbers,
    *,
    indices_are_sorted: bool = False,
    unique_indices: bool = False,
    mode: str | GatherScatterMode | None = None,
) -> Array:
  """Scatter-sub operator.

  Wraps `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_, where
  subtraction is used to combine updates and values from `operand`.

  The semantics of scatter are complicated, and its API might change in the
  future. For most use cases, you should prefer the
  :attr:`jax.numpy.ndarray.at` property on JAX arrays which uses
  the familiar NumPy indexing syntax.

  Args:
    operand: an array to which the scatter should be applied
    scatter_indices: an array that gives the indices in `operand` to which each
      update in `updates` should be applied.
    updates: the updates that should be scattered onto `operand`.
    dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes how
      dimensions of `operand`, `start_indices`, `updates` and the output relate.
    indices_are_sorted: whether `scatter_indices` is known to be sorted. If
      true, may improve performance on some backends.
    unique_indices: whether the elements to be updated in ``operand`` are
      guaranteed to not overlap with each other. If true, may improve
      performance on some backends. JAX does not check this promise: if the
      updated elements overlap when ``unique_indices`` is ``True`` the behavior
      is undefined.
    mode: how to handle indices that are out of bounds: when set to 'clip',
      indices are clamped so that the slice is within bounds, and when set to
      'fill' or 'drop' out-of-bounds updates are dropped. The behavior for
      out-of-bounds indices when set to 'promise_in_bounds' is
      implementation-defined.

  Returns:
    An array containing the difference between `operand` and the scattered
    updates.
  """
  jaxpr, consts = lax._reduction_jaxpr(
      lax.sub, core.get_aval(lax._const(operand, 0))
  )
  operand, scatter_indices, updates = core.standard_insert_pvary(
      operand, scatter_indices, updates)
  return scatter_sub_p.bind(
      operand,
      scatter_indices,
      updates,
      update_jaxpr=jaxpr,
      update_consts=consts,
      dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted,
      unique_indices=unique_indices,
      mode=GatherScatterMode.from_any(mode),
  )


@partial(api_boundary, repro_api_name="lax.scatter_mul")
def scatter_mul(
  operand: ArrayLike, scatter_indices: ArrayLike, updates: ArrayLike,
  dimension_numbers: ScatterDimensionNumbers, *,
  indices_are_sorted: bool = False, unique_indices: bool = False,
  mode: str | GatherScatterMode | None = None) -> Array:
  """Scatter-multiply operator.

  Wraps `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_, where
  multiplication is used to combine updates and values from `operand`.

  The semantics of scatter are complicated, and its API might change in the
  future. For most use cases, you should prefer the
  :attr:`jax.numpy.ndarray.at` property on JAX arrays which uses
  the familiar NumPy indexing syntax.

  Args:
    operand: an array to which the scatter should be applied
    scatter_indices: an array that gives the indices in `operand` to which each
      update in `updates` should be applied.
    updates: the updates that should be scattered onto `operand`.
    dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes
      how dimensions of `operand`, `start_indices`, `updates` and the output
      relate.
    indices_are_sorted: whether `scatter_indices` is known to be sorted. If
      true, may improve performance on some backends.
    unique_indices: whether the elements to be updated in ``operand`` are
      guaranteed to not overlap with each other. If true, may improve performance on
      some backends. JAX does not check this promise: if the updated elements
      overlap when ``unique_indices`` is ``True`` the behavior is undefined.
    mode: how to handle indices that are out of bounds: when set to 'clip',
      indices are clamped so that the slice is within bounds, and when
      set to 'fill' or 'drop' out-of-bounds updates are dropped. The behavior
      for out-of-bounds indices when set to 'promise_in_bounds' is
      implementation-defined.

  Returns:
    An array containing the product of `operand` and the scattered updates.
  """
  jaxpr, consts = lax._reduction_jaxpr(lax.mul,
                                       core.get_aval(lax._const(operand, 1)))
  operand, scatter_indices, updates = core.standard_insert_pvary(
      operand, scatter_indices, updates)
  return scatter_mul_p.bind(
      operand, scatter_indices, updates, update_jaxpr=jaxpr,
      update_consts=consts, dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=GatherScatterMode.from_any(mode))

@partial(api_boundary, repro_api_name="lax.scatter_min")
def scatter_min(
  operand: ArrayLike, scatter_indices: ArrayLike, updates: ArrayLike,
  dimension_numbers: ScatterDimensionNumbers, *,
  indices_are_sorted: bool = False, unique_indices: bool = False,
  mode: str | GatherScatterMode | None = None) -> Array:
  """Scatter-min operator.

  Wraps `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_, where
  the `min` function is used to combine updates and values from `operand`.

  The semantics of scatter are complicated, and its API might change in the
  future. For most use cases, you should prefer the
  :attr:`jax.numpy.ndarray.at` property on JAX arrays which uses
  the familiar NumPy indexing syntax.

  Args:
    operand: an array to which the scatter should be applied
    scatter_indices: an array that gives the indices in `operand` to which each
      update in `updates` should be applied.
    updates: the updates that should be scattered onto `operand`.
    dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes
      how dimensions of `operand`, `start_indices`, `updates` and the output
      relate.
    indices_are_sorted: whether `scatter_indices` is known to be sorted. If
      true, may improve performance on some backends.
    unique_indices: whether the elements to be updated in ``operand`` are
      guaranteed to not overlap with each other. If true, may improve performance on
      some backends. JAX does not check this promise: if the updated elements
      overlap when ``unique_indices`` is ``True`` the behavior is undefined.
    mode: how to handle indices that are out of bounds: when set to 'clip',
      indices are clamped so that the slice is within bounds, and when
      set to 'fill' or 'drop' out-of-bounds updates are dropped. The behavior
      for out-of-bounds indices when set to 'promise_in_bounds' is
      implementation-defined.

  Returns:
    An array containing the min of `operand` and the scattered updates.
  """
  jaxpr, consts = lax._reduction_jaxpr(lax.min,
                                       core.get_aval(lax._const(operand, 0)))
  operand, scatter_indices, updates = core.standard_insert_pvary(
      operand, scatter_indices, updates)
  return scatter_min_p.bind(
      operand, scatter_indices, updates, update_jaxpr=jaxpr,
      update_consts=consts, dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=GatherScatterMode.from_any(mode))

@partial(api_boundary, repro_api_name="lax.scatter_max")
def scatter_max(
  operand: ArrayLike, scatter_indices: ArrayLike, updates: ArrayLike,
  dimension_numbers: ScatterDimensionNumbers, *,
  indices_are_sorted: bool = False, unique_indices: bool = False,
  mode: str | GatherScatterMode | None = None) -> Array:
  """Scatter-max operator.

  Wraps `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_, where
  the `max` function is used to combine updates and values from `operand`.

  The semantics of scatter are complicated, and its API might change in the
  future. For most use cases, you should prefer the
  :attr:`jax.numpy.ndarray.at` property on JAX arrays which uses
  the familiar NumPy indexing syntax.

  Args:
    operand: an array to which the scatter should be applied
    scatter_indices: an array that gives the indices in `operand` to which each
      update in `updates` should be applied.
    updates: the updates that should be scattered onto `operand`.
    dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes
      how dimensions of `operand`, `start_indices`, `updates` and the output
      relate.
    indices_are_sorted: whether `scatter_indices` is known to be sorted. If
      true, may improve performance on some backends.
    unique_indices: whether the elements to be updated in ``operand`` are
      guaranteed to not overlap with each other. If true, may improve performance on
      some backends. JAX does not check this promise: if the updated elements
      overlap when ``unique_indices`` is ``True`` the behavior is undefined.
    mode: how to handle indices that are out of bounds: when set to 'clip',
      indices are clamped so that the slice is within bounds, and when
      set to 'fill' or 'drop' out-of-bounds updates are dropped. The behavior
      for out-of-bounds indices when set to 'promise_in_bounds' is
      implementation-defined.

  Returns:
    An array containing the max of `operand` and the scattered updates.
  """
  jaxpr, consts = lax._reduction_jaxpr(lax.max,
                                       core.get_aval(lax._const(operand, 0)))
  operand, scatter_indices, updates = core.standard_insert_pvary(
      operand, scatter_indices, updates)
  return scatter_max_p.bind(
      operand, scatter_indices, updates, update_jaxpr=jaxpr,
      update_consts=consts, dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=GatherScatterMode.from_any(mode))

# To avoid recompilation, we store a dict of weak references to funcs.
_scatter_apply_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()

def scatter_apply(
  operand: Array, scatter_indices: Array,
  func: Callable[[Array], Array],
  dimension_numbers: ScatterDimensionNumbers, *,
  update_shape: Shape = (),
  indices_are_sorted: bool = False, unique_indices: bool = False,
  mode: str | GatherScatterMode | None = None) -> Array:
  """Scatter-apply operator.

  Wraps `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_, where values
  from ``operand`` are replaced with ``func(operand)``, with duplicate indices
  resulting in multiple applications of ``func``.

  The semantics of scatter are complicated, and its API might change in the
  future. For most use cases, you should prefer the
  :attr:`jax.numpy.ndarray.at` property on JAX arrays which uses
  the familiar NumPy indexing syntax.

  Note that in the current implementation, ``scatter_apply`` is not compatible
  with automatic differentiation.

  Args:
    operand: an array to which the scatter should be applied
    scatter_indices: an array that gives the indices in `operand` to which each
      update in `updates` should be applied.
    func: unary function that will be applied at each index.
    dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes
      how dimensions of `operand`, `start_indices`, `updates` and the output
      relate.
    update_shape: the shape of the updates at the given indices.
    indices_are_sorted: whether `scatter_indices` is known to be sorted. If
      true, may improve performance on some backends.
    unique_indices: whether the elements to be updated in ``operand`` are
      guaranteed to not overlap with each other. If true, may improve performance on
      some backends. JAX does not check this promise: if the updated elements
      overlap when ``unique_indices`` is ``True`` the behavior is undefined.
    mode: how to handle indices that are out of bounds: when set to 'clip',
      indices are clamped so that the slice is within bounds, and when
      set to 'fill' or 'drop' out-of-bounds updates are dropped. The behavior
      for out-of-bounds indices when set to 'promise_in_bounds' is
      implementation-defined.

  Returns:
    An array containing the result of applying `func` to `operand` at the given indices.
  """
  # TODO: can we implement this without a placeholder?
  unused = lax.full(update_shape, 0, operand.dtype)
  _apply = lambda x, _: func(x)
  try:
    _apply = _scatter_apply_cache.setdefault(func, _apply)
  except TypeError:  # func is not weak referenceable
    pass
  jaxpr, consts = lax._reduction_jaxpr(_apply, core.get_aval(lax._zero(operand)))
  # TODO: implement this via its own primitive so we can define appropriate autodiff rules.
  operand, scatter_indices, unused = core.standard_insert_pvary(
      operand, scatter_indices, unused)
  return scatter_p.bind(
      operand, scatter_indices, unused, update_jaxpr=jaxpr,
      update_consts=consts, dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=GatherScatterMode.from_any(mode))

# Define this outside of scatter to ensure cache hits.
_scatter_reduction_computation = lambda x, y: y

def scatter(
  operand: ArrayLike, scatter_indices: ArrayLike, updates: ArrayLike,
  dimension_numbers: ScatterDimensionNumbers, *,
  indices_are_sorted: bool = False, unique_indices: bool = False,
  mode: str | GatherScatterMode | None = None) -> Array:
  """Scatter-update operator.

  Wraps `XLA's Scatter operator
  <https://www.openxla.org/xla/operation_semantics#scatter>`_, where updates
  replace values from `operand`.

  If multiple updates are performed to the same index of operand, they may be
  applied in any order.

  :func:`scatter` is a low-level operator with complicated semantics, and most
  JAX users will never need to call it directly. Instead, you should prefer using
  :func:`jax.numpy.ndarray.at` for more familiary NumPy-style indexing syntax.

  Args:
    operand: an array to which the scatter should be applied
    scatter_indices: an array that gives the indices in `operand` to which each
      update in `updates` should be applied.
    updates: the updates that should be scattered onto `operand`.
    dimension_numbers: a `lax.ScatterDimensionNumbers` object that describes
      how dimensions of `operand`, `start_indices`, `updates` and the output
      relate.
    indices_are_sorted: whether `scatter_indices` is known to be sorted. If
      true, may improve performance on some backends.
    unique_indices: whether the elements to be updated in ``operand`` are
      guaranteed to not overlap with each other. If true, may improve performance on
      some backends. JAX does not check this promise: if the updated elements
      overlap when ``unique_indices`` is ``True`` the behavior is undefined.
    mode: how to handle indices that are out of bounds: when set to 'clip',
      indices are clamped so that the slice is within bounds, and when
      set to 'fill' or 'drop' out-of-bounds updates are dropped. The behavior
      for out-of-bounds indices when set to 'promise_in_bounds' is
      implementation-defined.

  Returns:
    An array containing the values of `operand` and the scattered updates.

  Examples:
    As mentioned above, you should basically never use :func:`scatter` directly,
    and instead perform scatter-style operations using NumPy-style indexing
    expressions via :attr:`jax.numpy.ndarray.at`.

    Here is and example of updating entries in an array using :attr:`jax.numpy.ndarray.at`,
    which lowers to an XLA Scatter operation:

    >>> x = jnp.ones(5)
    >>> indices = jnp.array([1, 2, 4])
    >>> values = jnp.array([2.0, 3.0, 4.0])

    >>> x.at[indices].set(values)
    Array([1., 2., 3., 1., 4.], dtype=float32)

    This syntax also supports several of the optional arguments to :func:`scatter`,
    for example:

    >>> x.at[indices].set(values, indices_are_sorted=True, mode='promise_in_bounds')
    Array([1., 2., 3., 1., 4.], dtype=float32)

    By comparison, here is the equivalent function call using :func:`scatter` directly,
    which is not something typical users should ever need to do:

    >>> lax.scatter(x, indices[:, None], values,
    ...             dimension_numbers=lax.ScatterDimensionNumbers(
    ...                 update_window_dims=(),
    ...                 inserted_window_dims=(0,),
    ...                 scatter_dims_to_operand_dims=(0,)),
    ...             indices_are_sorted=True,
    ...             mode=lax.GatherScatterMode.PROMISE_IN_BOUNDS)
    Array([1., 2., 3., 1., 4.], dtype=float32)
  """
  operand, scatter_indices, updates = core.standard_insert_pvary(
      operand, scatter_indices, updates)
  return scatter_p.bind(
      operand, scatter_indices, updates, update_jaxpr=None,
      update_consts=(), dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=GatherScatterMode.from_any(mode))

def index_take(src: Array, idxs: Array, axes: Sequence[int]) -> Array:
  indices = lax.concatenate([lax.expand_dims(i, (1,)) for i in idxs], 1)
  max_idx = lax.expand_dims(np.array([src.shape[ax] for ax in axes]),
                            tuple(range(indices.ndim - 1)))
  indices = indices % max_idx
  slice_sizes = list(src.shape)
  for ax in axes:
    slice_sizes[ax] = 1
  offset_dims = tuple(range(1, src.ndim - indices.shape[1] + 1))
  dnums = GatherDimensionNumbers(
      offset_dims=offset_dims,
      collapsed_slice_dims=tuple(axes),
      start_index_map=tuple(axes),
  )
  return gather(src, indices, dimension_numbers=dnums,
                slice_sizes=tuple(slice_sizes))


### convenience wrappers around traceables

def slice_in_dim(operand: Array | np.ndarray, start_index: int | None,
                 limit_index: int | None,
                 stride: int = 1, axis: int = 0) -> Array:
  """Convenience wrapper around :func:`lax.slice` applying to only one dimension.

  This is effectively equivalent to ``operand[..., start_index:limit_index:stride]``
  with the indexing applied on the specified axis.

  Args:
    operand: an array to slice.
    start_index: an optional start index (defaults to zero)
    limit_index: an optional end index (defaults to operand.shape[axis])
    stride: an optional stride (defaults to 1)
    axis: the axis along which to apply the slice (defaults to 0)

  Returns:
    An array containing the slice.

  Examples:
    Here is a one-dimensional example:

    >>> x = jnp.arange(4)
    >>> lax.slice_in_dim(x, 1, 3)
    Array([1, 2], dtype=int32)

    Here are some two-dimensional examples:

    >>> x = jnp.arange(12).reshape(4, 3)
    >>> x
    Array([[ 0,  1,  2],
           [ 3,  4,  5],
           [ 6,  7,  8],
           [ 9, 10, 11]], dtype=int32)

    >>> lax.slice_in_dim(x, 1, 3)
    Array([[3, 4, 5],
           [6, 7, 8]], dtype=int32)

    >>> lax.slice_in_dim(x, 1, 3, axis=1)
    Array([[ 1,  2],
           [ 4,  5],
           [ 7,  8],
           [10, 11]], dtype=int32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :func:`jax.lax.slice`
    - :func:`jax.lax.index_in_dim`
    - :func:`jax.lax.dynamic_slice_in_dim`
  """
  start_indices = [0] * operand.ndim
  limit_indices = list(operand.shape)
  strides = [1] * operand.ndim

  # translate `None`
  len_axis = operand.shape[axis]
  start_index_int = (core._canonicalize_dimension(start_index)
                     if start_index is not None else 0)
  limit_index_int = (core._canonicalize_dimension(limit_index)
                     if limit_index is not None else len_axis)

  # translate negative indices
  if start_index_int < 0:
    start_index_int = start_index_int + len_axis
  if limit_index_int < 0:
    limit_index_int = limit_index_int + len_axis

  axis = int(axis)
  start_indices[axis] = start_index_int
  limit_indices[axis] = limit_index_int
  strides[axis] = core._canonicalize_dimension(stride)

  return slice(operand, start_indices, limit_indices, strides)


def index_in_dim(operand: Array | np.ndarray, index: int, axis: int = 0,
                 keepdims: bool = True) -> Array:
  """Convenience wrapper around :func:`lax.slice` to perform int indexing.

  This is effectively equivalent to ``operand[..., index]``
  with the indexing applied on the specified axis.

  Args:
    operand: an array to index.
    index: integer index
    axis: the axis along which to apply the index (defaults to 0)
    keepdims: boolean specifying whether the output array should preserve the
      rank of the input (default=True)

  Returns:
    The subarray at the specified index.

  Examples:
    Here is a one-dimensional example:

    >>> x = jnp.arange(4)
    >>> lax.index_in_dim(x, 2)
    Array([2], dtype=int32)

    >>> lax.index_in_dim(x, 2, keepdims=False)
    Array(2, dtype=int32)

    Here are some two-dimensional examples:

    >>> x = jnp.arange(12).reshape(3, 4)
    >>> x
    Array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]], dtype=int32)

    >>> lax.index_in_dim(x, 1)
    Array([[4, 5, 6, 7]], dtype=int32)

    >>> lax.index_in_dim(x, 1, axis=1, keepdims=False)
    Array([1, 5, 9], dtype=int32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :func:`jax.lax.slice`
    - :func:`jax.lax.slice_in_dim`
    - :func:`jax.lax.dynamic_index_in_dim`
  """
  index, axis = core._canonicalize_dimension(index), int(axis)
  axis_size = operand.shape[axis]
  wrapped_index = index + axis_size if index < 0 else index
  if not 0 <= wrapped_index < axis_size:
    msg = 'index {} is out of bounds for axis {} with size {}'
    raise IndexError(msg.format(index, axis, axis_size))
  result = slice_in_dim(operand, wrapped_index, wrapped_index + 1, 1, axis)
  if keepdims:
    return result
  else:
    return lax.squeeze(result, (axis,))


def dynamic_slice_in_dim(operand: Array | np.ndarray,
                         start_index: ArrayLike,
                         slice_size: int, axis: int = 0, *,
                         allow_negative_indices: bool = True) -> Array:
  """Convenience wrapper around :func:`lax.dynamic_slice` applied to one dimension.

  This is roughly equivalent to the following Python indexing syntax applied
  along the specified axis: ``operand[..., start_index:start_index + slice_size]``.

  Args:
    operand: an array to slice.
    start_index: the (possibly dynamic) start index
    slice_size: the static slice size
    axis: the axis along which to apply the slice (defaults to 0)
    allow_negative_indices: boolean specifying whether negative indices are
      allowed. If true, negative indices are taken relative to the end of the
      array. If false, negative indices are out of bounds and the result is
      implementation defined.

  Returns:
    An array containing the slice.

  Examples:
    Here is a one-dimensional example:

    >>> x = jnp.arange(5)
    >>> dynamic_slice_in_dim(x, 1, 3)
    Array([1, 2, 3], dtype=int32)

    Like `jax.lax.dynamic_slice`, out-of-bound slices will be clipped to the
    valid range:

    >>> dynamic_slice_in_dim(x, 4, 3)
    Array([2, 3, 4], dtype=int32)

    Here is a two-dimensional example:

    >>> x = jnp.arange(12).reshape(3, 4)
    >>> x
    Array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]], dtype=int32)

    >>> dynamic_slice_in_dim(x, 1, 2, axis=1)
    Array([[ 1,  2],
           [ 5,  6],
           [ 9, 10]], dtype=int32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :func:`jax.lax.slice_in_dim`
    - :func:`jax.lax.dynamic_slice`
    - :func:`jax.lax.dynamic_index_in_dim`
  """
  start_indices: list[ArrayLike] = [lax._const(start_index, 0)] * operand.ndim
  slice_sizes = list(operand.shape)
  axis = int(axis)
  allow_negative = [False] * operand.ndim
  allow_negative[axis] = allow_negative_indices
  start_indices[axis] = start_index
  slice_sizes[axis] = core._canonicalize_dimension(slice_size)
  return dynamic_slice(operand, start_indices, slice_sizes,
                       allow_negative_indices=allow_negative)


def dynamic_index_in_dim(operand: Array | np.ndarray,
                         index: ArrayLike,
                         axis: int = 0, keepdims: bool = True,
                         *,
                         allow_negative_indices: bool = True) -> Array:
  """Convenience wrapper around dynamic_slice to perform int indexing.

  This is roughly equivalent to the following Python indexing syntax applied
  along the specified axis: ``operand[..., index]``.

  Args:
    operand: an array to slice.
    index: the (possibly dynamic) start index
    axis: the axis along which to apply the slice (defaults to 0)
    keepdims: boolean specifying whether the output should have the same rank as
      the input (default = True)
    allow_negative_indices: boolean specifying whether negative indices are
      allowed. If true, negative indices are taken relative to the end of the
      array. If false, negative indices are out of bounds and the result is
      implementation defined.

  Returns:
    An array containing the slice.

  Examples:
    Here is a one-dimensional example:

    >>> x = jnp.arange(5)
    >>> dynamic_index_in_dim(x, 1)
    Array([1], dtype=int32)

    >>> dynamic_index_in_dim(x, 1, keepdims=False)
    Array(1, dtype=int32)

    Here is a two-dimensional example:

    >>> x = jnp.arange(12).reshape(3, 4)
    >>> x
    Array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]], dtype=int32)

    >>> dynamic_index_in_dim(x, 1, axis=1, keepdims=False)
    Array([1, 5, 9], dtype=int32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :func:`jax.lax.index_in_dim`
    - :func:`jax.lax.dynamic_slice`
    - :func:`jax.lax.dynamic_slice_in_dim`
  """
  result = dynamic_slice_in_dim(operand, index, 1, axis,
                                allow_negative_indices=allow_negative_indices)
  if keepdims:
    return result
  else:
    return lax.squeeze(result, (axis,))


def dynamic_update_slice_in_dim(operand: Array | np.ndarray,
                                update: ArrayLike,
                                start_index: ArrayLike, axis: int,
                                *,
                                allow_negative_indices: bool = True) -> Array:
  """Convenience wrapper around :func:`dynamic_update_slice` to update
  a slice in a single ``axis``.

  Args:
    operand: an array to slice.
    update: an array containing the new values to write onto `operand`.
    start_index: a single scalar index
    axis: the axis of the update.
    allow_negative_indices: boolean specifying whether negative indices are
      allowed. If true, negative indices are taken relative to the end of the
      array. If false, negative indices are out of bounds and the result is
      implementation defined.

  Returns:
    The updated array

  Examples:

    >>> x = jnp.zeros(6)
    >>> y = jnp.ones(3)
    >>> dynamic_update_slice_in_dim(x, y, 2, axis=0)
    Array([0., 0., 1., 1., 1., 0.], dtype=float32)

    If the update slice is too large to fit in the array, the start
    index will be adjusted to make it fit:

    >>> dynamic_update_slice_in_dim(x, y, 3, axis=0)
    Array([0., 0., 0., 1., 1., 1.], dtype=float32)
    >>> dynamic_update_slice_in_dim(x, y, 5, axis=0)
    Array([0., 0., 0., 1., 1., 1.], dtype=float32)

    Here is an example of a two-dimensional slice update:

    >>> x = jnp.zeros((4, 4))
    >>> y = jnp.ones((2, 4))
    >>> dynamic_update_slice_in_dim(x, y, 1, axis=0)
    Array([[0., 0., 0., 0.],
           [1., 1., 1., 1.],
           [1., 1., 1., 1.],
           [0., 0., 0., 0.]], dtype=float32)

    Note that the shape of the additional axes in ``update`` need not
    match the associated dimensions of the ``operand``:

    >>> y = jnp.ones((2, 3))
    >>> dynamic_update_slice_in_dim(x, y, 1, axis=0)
    Array([[0., 0., 0., 0.],
           [1., 1., 1., 0.],
           [1., 1., 1., 0.],
           [0., 0., 0., 0.]], dtype=float32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :func:`jax.lax.dynamic_update_slice`
    - :func:`jax.lax.dynamic_update_index_in_dim`
    - :func:`jax.lax.dynamic_slice_in_dim`
  """
  axis = int(axis)
  start_indices: list[ArrayLike] = [lax._const(start_index, 0)] * lax._ndim(operand)
  start_indices[axis] = start_index
  allow_negative = [False] * operand.ndim
  allow_negative[axis] = allow_negative_indices
  return dynamic_update_slice(operand, update, start_indices,
                              allow_negative_indices=allow_negative)


def dynamic_update_index_in_dim(operand: Array | np.ndarray,
                                update: ArrayLike, index: ArrayLike,
                                axis: int, *,
                                allow_negative_indices: bool = True) -> Array:
  """Convenience wrapper around :func:`dynamic_update_slice` to update a slice
  of size 1 in a single ``axis``.

  Args:
    operand: an array to slice.
    update: an array containing the new values to write onto `operand`.
    index: a single scalar index
    axis: the axis of the update.
    allow_negative_indices: boolean specifying whether negative indices are
      allowed. If true, negative indices are taken relative to the end of the
      array. If false, negative indices are out of bounds and the result is
      implementation defined.

  Returns:
    The updated array

  Examples:

    >>> x = jnp.zeros(6)
    >>> y = 1.0
    >>> dynamic_update_index_in_dim(x, y, 2, axis=0)
    Array([0., 0., 1., 0., 0., 0.], dtype=float32)

    >>> y = jnp.array([1.0])
    >>> dynamic_update_index_in_dim(x, y, 2, axis=0)
    Array([0., 0., 1., 0., 0., 0.], dtype=float32)

    If the specified index is out of bounds, the index will be clipped to the
    valid range:

    >>> dynamic_update_index_in_dim(x, y, 10, axis=0)
    Array([0., 0., 0., 0., 0., 1.], dtype=float32)

    Here is an example of a two-dimensional dynamic index update:

    >>> x = jnp.zeros((4, 4))
    >>> y = jnp.ones(4)
    >>> dynamic_update_index_in_dim(x, y, 1, axis=0)
    Array([[0., 0., 0., 0.],
          [1., 1., 1., 1.],
          [0., 0., 0., 0.],
          [0., 0., 0., 0.]], dtype=float32)

    Note that the shape of the additional axes in ``update`` need not
    match the associated dimensions of the ``operand``:

    >>> y = jnp.ones((1, 3))
    >>> dynamic_update_index_in_dim(x, y, 1, 0)
    Array([[0., 0., 0., 0.],
           [1., 1., 1., 0.],
           [0., 0., 0., 0.],
           [0., 0., 0., 0.]], dtype=float32)

  See Also:
    - :attr:`jax.numpy.ndarray.at`
    - :func:`jax.lax.dynamic_update_slice`
    - :func:`jax.lax.dynamic_update_index_in_dim`
    - :func:`jax.lax.dynamic_index_in_dim`
  """
  axis = int(axis)
  if lax._ndim(update) != lax._ndim(operand):
    assert lax._ndim(update) + 1 == lax._ndim(operand)
    update = lax.expand_dims(update, (axis,))
  return dynamic_update_slice_in_dim(
      operand, update, index, axis,
      allow_negative_indices=allow_negative_indices)


def _slice_shape_rule(operand, *, start_indices, limit_indices, strides):
  lax._check_shapelike("slice", "start_indices", start_indices)
  lax._check_shapelike("slice", "limit_indices", limit_indices)
  if operand.ndim != len(start_indices):
    msg = ("slice start_indices must have length equal to the number of "
           "dimensions of the operand, got indices {} for operand shape {}.")
    raise TypeError(msg.format(start_indices, operand.shape))
  if len(start_indices) != len(limit_indices):
    msg = ("slice limit_indices must have the same length as start_indices, "
           "got start_indices {} and limit_indices {}.")
    raise TypeError(msg.format(start_indices, limit_indices))
  if not all(map(operator.ge, operand.shape, limit_indices)):
    msg = ("slice limit_indices must be less than or equal to operand shape, "
           "got limit_indices {} for operand shape {}.")
    raise TypeError(msg.format(limit_indices, operand.shape))
  if not all(si >= 0 for si in start_indices):
    msg = ("slice start_indices must be greater than or equal to zero, "
           "got start_indices of {}.")
    raise TypeError(msg.format(start_indices))
  if not all(map(operator.ge, limit_indices, start_indices)):
    msg = ("slice limit_indices must be greater than or equal to start_indices,"
          " got start_indices {} and limit_indices {}.")
    raise TypeError(msg.format(start_indices, limit_indices))
  diff = tuple(map(operator.sub, limit_indices, start_indices))
  if strides is None or tuple(strides) == (1,) * len(operand.shape):
    return diff

  lax._check_shapelike("slice", "strides", strides)
  if len(strides) != operand.ndim:
    msg = ("slice strides must have length equal to the number of dimensions "
            "of the operand, got strides {} for operand shape {}.")
    raise TypeError(msg.format(strides, operand.shape))
  if not all(s >= 0 for s in strides):
    msg = "slice strides must be positive, got {}"
    raise TypeError(msg.format(strides))
  return tuple(core.stride_dim(d, window_size=1, window_stride=s)
               for d, s in zip(diff, strides))

def _get_sub_spec_size(mesh, sub_spec):
  if isinstance(sub_spec, tuple):
    return math.prod(mesh.shape[s] for s in sub_spec)
  return mesh.shape[sub_spec]

def _get_sharding_for_varying_out_shape(out_shape, operand, name):
  """Returns a sharding when out_shape may not be the same as operand shape"""
  mesh = operand.sharding.mesh
  for op_sh, out_sh, op_spec in safe_zip(
      operand.shape, out_shape, operand.sharding.spec):
    if (op_sh != out_sh and op_spec is not None and
        out_sh % _get_sub_spec_size(mesh, op_spec) != 0):
      raise core.ShardingTypeError(
          f"{name} on sharded dims where out dim ({out_sh}) is not divisible by"
          f" mesh axes ({_get_sub_spec_size(mesh, op_spec)}) with spec"
          f" ({op_spec}) is not implemented."
      )
  # TODO(yashkatariya): Returning operand.sharding as is may or may not move
  # data. So think about how to avoid it which might include creating a new
  # mesh? For example:
  # mesh = {'x': 4}
  # x = jax.device_put(jnp.arange(8), NamedSharding(mesh, P('x')))`
  # ys = lax.split(x, [4, 4])  # This will create outputs of shape (4,)
  # According to the current logic, ys[0].sharding.spec == P('x')
  # which involves data movement.
  return operand.sharding

def _slice_sharding_rule(operand, *, start_indices, limit_indices, strides):
  # TODO(yashkatariya): Once JAX supports uneven sharding at the top level,
  # change this logic to `return operand.sharding` directly.
  out_shape = _slice_shape_rule(operand, start_indices=start_indices,
                                limit_indices=limit_indices, strides=strides)
  return _get_sharding_for_varying_out_shape(out_shape, operand, 'slicing')

def _slice_transpose_rule(t, operand, *, start_indices, limit_indices, strides):
  operand_shape = operand.aval.shape
  if strides is None or np.all(np.equal(strides, 1)):
    pads = zip(start_indices, np.subtract(operand_shape, limit_indices),
               (0,) * len(start_indices))
  else:
    real_limits = np.add(
      start_indices,
      np.where(np.array(t.shape) == 0, 0,
               np.add(1, np.multiply(np.subtract(t.shape, 1), strides))))
    pads = zip(start_indices, np.subtract(operand_shape, real_limits),
               np.subtract(strides, 1))
  result = lax.pad(t, lax._const(t, 0), pads)
  assert result.shape == operand_shape, f"{result.shape=} {operand_shape=}"
  return [result]

def _slice_transpose_fancy(out_ct, operand, *, start_indices, limit_indices, strides):
  assert isinstance(operand, ad.GradAccum)
  if type(out_ct) is ad_util.Zero: return
  if isinstance(operand, ad.RefAccum):
    slices = map(_slice, start_indices, limit_indices, strides)
    operand.ref.addupdate(out_ct, tuple(slices))
  else:
    if strides is None or np.all(np.equal(strides, 1)):
      pads = zip(start_indices, np.subtract(operand.aval.shape, limit_indices),
                 (0,) * len(start_indices))
    else:
      real_limits = np.add(
        start_indices,
        np.where(np.array(out_ct.shape) == 0, 0,
                 np.add(1, np.multiply(np.subtract(out_ct.shape, 1), strides))))
      pads = zip(start_indices, np.subtract(operand.aval.shape, real_limits),
                 np.subtract(strides, 1))
    operand.accum(lax.pad(out_ct, lax._const(out_ct, 0), pads))


def _slice_batching_rule(batched_args, batch_dims, *, start_indices,
                         limit_indices, strides):
  operand, = batched_args
  bdim, = batch_dims

  new_start_indices = list(start_indices)
  new_start_indices.insert(bdim, 0)

  new_limit_indices = list(limit_indices)
  new_limit_indices.insert(bdim, operand.shape[bdim])

  if strides is None:
    new_strides = None
  else:
    new_strides = list(strides)
    new_strides.insert(bdim, 1)

  out = slice(operand, new_start_indices, new_limit_indices, new_strides)
  return out, bdim

slice_p = standard_primitive(_slice_shape_rule, input_dtype, 'slice',
                             sharding_rule=_slice_sharding_rule,
                             vma_rule=partial(core.standard_vma_rule, 'slice'))
ad.deflinear2(slice_p, _slice_transpose_rule)
ad.fancy_transposes[slice_p] = _slice_transpose_fancy
batching.primitive_batchers[slice_p] = _slice_batching_rule

# Override the standard impl to defer to dynamic_slice whenever possible.
# This lets us reuse the same program for many applications of slicing for as
# long as they have the same output shape. Note that this only applies in
# op-by-op mode and inside larger computations we still lower to static slices.
@slice_p.def_impl
def _slice_impl(x, start_indices, limit_indices, strides):
  if strides is not None:
    return dispatch.apply_primitive(
      slice_p, x, start_indices=start_indices,
      limit_indices=limit_indices, strides=strides)
  slice_sizes = tuple(np.array(limit_indices) - np.array(start_indices))
  return dispatch.apply_primitive(dynamic_slice_p, x, *start_indices,
                                  slice_sizes=slice_sizes)


def _slice_lower(ctx, x, *, start_indices, limit_indices, strides):
  strides = strides or [1] * len(start_indices)
  aval_out, = ctx.avals_out
  out = mlir.slice_op(ctx, x, aval_out, start_indices=start_indices,
                      limit_indices=limit_indices, strides=strides)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

mlir.register_lowering(slice_p, _slice_lower)


def _dynamic_slice_shape_rule(operand, *start_indices, slice_sizes):
  if not all(map(operator.ge, operand.shape, slice_sizes)):
    msg = ("slice slice_sizes must be less than or equal to operand shape, "
           "got slice_sizes {} for operand shape {}.")
    raise TypeError(msg.format(slice_sizes, operand.shape))
  if not all(ssz >= 0 for ssz in slice_sizes):
    msg = ("slice slice_sizes must be greater than or equal to zero, "
           "got slice_sizes of {}.")
    raise TypeError(msg.format(slice_sizes))
  if any(idx.ndim != 0 for idx in start_indices):
    raise TypeError("start_indices arguments to dynamic_slice must be scalars, "
                    f" got indices {start_indices}")
  return tuple(slice_sizes)

def _dynamic_slice_sharding_rule(operand, *starts_and_dyn_sizes, slice_sizes):
  out_shape = _dynamic_slice_shape_rule(
      operand, *starts_and_dyn_sizes, slice_sizes=slice_sizes)
  return _get_sharding_for_varying_out_shape(out_shape, operand, 'dynamic_slice')

def _dynamic_slice_reduced_rule(out_s, operand, *starts_and_dyn_sizes,
                                slice_sizes):
  return out_s.update(spec=out_s.spec.update(
      reduced=operand.sharding.spec.reduced))

def _dynamic_slice_dtype_rule(operand, *starts_and_dyn_sizes, slice_sizes):
  start_indices, dyn = util.split_list(starts_and_dyn_sizes, [operand.ndim])
  if any(i.dtype != start_indices[0].dtype or
         not dtypes.issubdtype(i.dtype, np.integer) for i in start_indices):
    msg = ("index arguments to dynamic_slice must be integers of the same "
           "type, got: {}")
    raise TypeError(msg.format(", ".join(i.dtype.name for i in start_indices)))
  return operand.dtype

def _dynamic_slice_jvp(primals, tangents, *, slice_sizes):
  tangent_out = tangents[0]
  if type(tangent_out) is not ad_util.Zero:
    tangent_out = dynamic_slice_p.bind(tangent_out, *primals[1:], slice_sizes=slice_sizes)
  return dynamic_slice_p.bind(primals[0], *primals[1:], slice_sizes=slice_sizes), tangent_out

def _dynamic_slice_transpose_rule(t, operand, *start_indices, slice_sizes):
  assert ad.is_undefined_primal(operand)
  assert all(not ad.is_undefined_primal(s) for s in start_indices)
  if type(t) is ad_util.Zero:
    return [ad_util.Zero(operand.aval)] + [None] * len(start_indices)
  else:
    zeros = lax.full(operand.aval.shape, 0, operand.aval.dtype,
                     sharding=operand.aval.sharding)
    zeros = core.pvary(zeros, tuple(operand.aval.vma))
    return ([dynamic_update_slice_p.bind(zeros, t, *start_indices)] +
            [None] * len(start_indices))

def _dynamic_slice_transpose_fancy(out_ct, operand, *start_indices, slice_sizes):
  assert isinstance(operand, ad.GradAccum)
  assert all(not isinstance(s, ad.GradAccum) for s in start_indices)
  if type(out_ct) is ad_util.Zero: return
  if isinstance(operand, ad.RefAccum):
    operand.ref.addupdate(out_ct, tuple(map(ds, start_indices, slice_sizes)))
  else:
    zeros = lax.full(operand.aval.shape, 0, operand.aval.dtype,
                     sharding=operand.aval.sharding)
    zeros = core.pvary(zeros, tuple(operand.aval.vma))
    operand.accum(dynamic_update_slice_p.bind(zeros, out_ct, *start_indices))

def _batch_dynamic_slice_indices(indices, bdims):
  if len(indices) == 0:
    return np.array([], 'int32'), None
  empty_marker = object()
  size = next((x.shape[i] for x, i in zip(indices, bdims) if i is not None),
              empty_marker)
  out = next(((core.typeof(x).sharding.mesh, core.typeof(x).sharding.spec[i])
              for x, i in zip(indices, bdims) if i is not None), None)
  if size is empty_marker:
    return lax.concatenate([lax.broadcast(i, (1,)) for i in indices], 0), None
  out_s = None if out is None else NamedSharding(out[0], P(out[1], None))
  indices = lax.concatenate(
    [lax.broadcast_in_dim(
        x, (size, 1), broadcast_dimensions=((0,) if i is not None else ()),
        out_sharding=out_s)
     for x, i in zip(indices, bdims)], dimension=1)
  return indices, 0

def _dynamic_slice_batching_rule(batched_args, batch_dims, *, slice_sizes):
  # A dynamic slice is a special case of gather; we can delegate to the gather
  # batching rule.
  # TODO(phawkins): consider removing dynamic_slice entirely and using gather
  # always.
  # TODO(mattjj): Alternately, we could add jnp.unstack and an unstack_p,
  # since it should have easier rules (especially compared to gather).
  operand, *start_indices_and_dyn = batched_args
  operand_bd, *start_idx_and_dyn_bds = batch_dims
  ndims = operand.ndim - (0 if operand_bd is batching.not_mapped else 1)
  dims = tuple(range(ndims))
  start_indices, dyn_slice_sizes = util.split_list(start_indices_and_dyn, [ndims])
  start_idx_bds, dyn_slice_size_bds = util.split_list(start_idx_and_dyn_bds, [ndims])
  dnums = GatherDimensionNumbers(
      offset_dims=dims,
      collapsed_slice_dims=(),
      start_index_map=dims,
  )
  index, index_bdim = _batch_dynamic_slice_indices(start_indices, start_idx_bds)
  return _gather_batching_rule(
    [operand, index, *dyn_slice_sizes],
    [operand_bd, index_bdim, *dyn_slice_size_bds], dimension_numbers=dnums,
    slice_sizes=slice_sizes, unique_indices=True, indices_are_sorted=True,
    mode=GatherScatterMode.CLIP, fill_value=None)

def _dynamic_slice_staging_rule(trace, source_info, x, *start_indices,
                                slice_sizes):
  return trace.default_process_primitive(
      dynamic_slice_p, (x, *start_indices), dict(slice_sizes=slice_sizes),
      source_info=source_info)

def _dynamic_slice_typecheck_rule(_, x, *start_indices, slice_sizes):
  out_aval, effects = dynamic_slice_p.abstract_eval(
      x.aval, *(d.aval for d in start_indices), slice_sizes=slice_sizes)
  return [out_aval], effects

dynamic_slice_p = standard_primitive(
    _dynamic_slice_shape_rule, _dynamic_slice_dtype_rule, 'dynamic_slice',
    weak_type_rule=_argnum_weak_type(0),
    sharding_rule=_dynamic_slice_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'dynamic_slice'),
    reduced_rule=_dynamic_slice_reduced_rule)
ad.primitive_jvps[dynamic_slice_p] = _dynamic_slice_jvp
ad.primitive_transposes[dynamic_slice_p] = _dynamic_slice_transpose_rule
ad.fancy_transposes[dynamic_slice_p] = _dynamic_slice_transpose_fancy
batching.primitive_batchers[dynamic_slice_p] = _dynamic_slice_batching_rule
pe.custom_staging_rules[dynamic_slice_p] = _dynamic_slice_staging_rule
core.custom_typechecks[dynamic_slice_p] = _dynamic_slice_typecheck_rule

def _dynamic_slice_lower(ctx, x, *start_indices, slice_sizes):
  x_aval, *_ = ctx.avals_in
  aval_out, = ctx.avals_out
  out = mlir.dynamic_slice(ctx, aval_out, x, start_indices=start_indices)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

mlir.register_lowering(dynamic_slice_p, _dynamic_slice_lower)


def _dynamic_update_slice_shape_rule(operand, update, *start_indices):
  if operand.ndim != update.ndim:
    msg = ("dynamic_update_slice update must have the same rank as operand, "
           "got update shape {} for operand shape {}.")
    raise TypeError(msg.format(update.shape, operand.shape))
  if operand.ndim != len(start_indices):
    msg = ("dynamic_update_slice start_indices must have length equal to the "
           "rank of operand, got indices {} for operand shape {}.")
    raise TypeError(msg.format(start_indices, operand.shape))
  if not all(map(operator.ge, operand.shape, update.shape)):
    msg = ("dynamic_update_slice update shape must be smaller than operand "
           "shape, got update shape {} for operand shape {}.")
    raise TypeError(msg.format(update.shape, operand.shape))
  if any(idx.ndim != 0 for idx in start_indices):
    raise TypeError("start_indices arguments to dynamic_update_slice must be "
                    f"scalars, got indices {start_indices}")
  return operand.shape

def _dynamic_update_slice_sharding_rule(operand, update, *start_indices):
  if operand.sharding != update.sharding:
    raise core.ShardingTypeError(
        "dynamic_update_slice operand sharding must be equal to update"
        " sharding, got operand sharding"
        f" {operand.str_short(mesh_axis_types=True)} and update sharding"
        f" {update.str_short(mesh_axis_types=True)}.")
  return operand.sharding

def _dynamic_update_slice_unreduced_rule(out_s, operand, update, *start_indices):
  if operand.sharding.spec.unreduced != update.sharding.spec.unreduced:
    raise core.ShardingTypeError(
        "dynamic_update_slice operand and update must be unreduced along the"
        " same axes. Got operand sharding"
        f" {operand.str_short(mesh_axis_types=True)} and update sharding"
        f" {update.str_short(mesh_axis_types=True)}.")
  return out_s.update(spec=out_s.spec.update(
      unreduced=operand.sharding.spec.unreduced))

def _dynamic_update_slice_reduced_rule(out_s, operand, update, *start_indices):
  if operand.sharding.spec.reduced != update.sharding.spec.reduced:
    raise core.ShardingTypeError(
        "dynamic_update_slice operand and update must be reduced along the"
        " same axes. Got operand sharding"
        f" {operand.str_short(mesh_axis_types=True)} and update sharding"
        f" {update.str_short(mesh_axis_types=True)}.")
  return out_s.update(spec=out_s.spec.update(
      reduced=operand.sharding.spec.reduced))

def _dynamic_update_slice_dtype_rule(operand, update, *start_indices):
  lax.check_same_dtypes("dynamic_update_slice", operand, update)
  if any(i.dtype != start_indices[0].dtype or
         not dtypes.issubdtype(i.dtype, np.integer) for i in start_indices):
    msg = ("index arguments to dynamic_update_slice must be integers of the "
           "same type, got {}")
    raise TypeError(msg.format(", ".join(i.dtype.name for i in start_indices)))
  return operand.dtype

def _dynamic_update_slice_jvp(primals, tangents):
  operand, update = primals[:2]
  start_indices = primals[2:]
  g_operand, g_update = tangents[:2]
  val_out = dynamic_update_slice_p.bind(operand, update, *start_indices)
  if type(g_operand) is ad_util.Zero and type(g_update) is ad_util.Zero:
    tangent_out = ad_util.Zero.from_primal_value(val_out)
  else:
    g_operand = ad.instantiate_zeros(g_operand)
    g_update = ad.instantiate_zeros(g_update)
    tangent_out = dynamic_update_slice_p.bind(g_operand, g_update, *start_indices)
  return val_out, tangent_out

def _dynamic_update_slice_transpose_rule(t, operand, update, *start_indices):
  assert all(not ad.is_undefined_primal(x) for x in start_indices)
  update_shape = (update.aval.shape if ad.is_undefined_primal(update) else
                  update.shape)
  update_sharding = update.aval.sharding
  if type(t) is ad_util.Zero:
    operand_t = ad_util.Zero(operand.aval) if ad.is_undefined_primal(operand) else None
    update_t = ad_util.Zero(update.aval) if ad.is_undefined_primal(update) else None
  else:
    zeros = lax._zeros(t, shape=update_shape, sharding=update_sharding)
    operand_t = (dynamic_update_slice_p.bind(t, zeros, *start_indices)
                 if ad.is_undefined_primal(operand) else None)
    update_t = (dynamic_slice_p.bind(t, *start_indices, slice_sizes=update_shape)
                if ad.is_undefined_primal(update) else None)
  return [operand_t, update_t] + [None] * len(start_indices)

def _dynamic_update_slice_batching_rule(batched_args, batch_dims):
  # A dynamic update slice is a special case of scatter; we can delegate to the
  # scatter batching rule.
  # TODO(phawkins): consider removing dynamic_update_slice entirely and using
  # scatter always.
  operand, update, *start_idx = batched_args
  operand_bd, update_bd, *start_idx_bd = batch_dims
  update_shape = (np.shape(update) if update_bd is batching.not_mapped
                  else tuple(np.delete(np.shape(update), update_bd)))
  dims = tuple(range(len(update_shape)))
  dnums = ScatterDimensionNumbers(
      update_window_dims=dims,
      inserted_window_dims=(),
      scatter_dims_to_operand_dims=dims,
  )
  index, index_bdim = _batch_dynamic_slice_indices(start_idx, start_idx_bd)
  return api.vmap(
    partial(scatter, dimension_numbers=dnums,
            indices_are_sorted=True, unique_indices=True,
            mode=GatherScatterMode.CLIP),
    in_axes=(operand_bd, index_bdim, update_bd),
    out_axes=0)(operand, index, update), 0


dynamic_update_slice_p = standard_primitive(
    _dynamic_update_slice_shape_rule, _dynamic_update_slice_dtype_rule,
    'dynamic_update_slice', sharding_rule=_dynamic_update_slice_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'dynamic_update_slice'),
    unreduced_rule=_dynamic_update_slice_unreduced_rule,
    reduced_rule=_dynamic_update_slice_reduced_rule)
ad.primitive_jvps[dynamic_update_slice_p] = _dynamic_update_slice_jvp
ad.primitive_transposes[dynamic_update_slice_p] = \
    _dynamic_update_slice_transpose_rule
batching.primitive_batchers[dynamic_update_slice_p] = \
    _dynamic_update_slice_batching_rule

def _dynamic_update_slice_lower(ctx, x, update, *start_indices):
  aval_out, = ctx.avals_out
  out = mlir.dynamic_update_slice(ctx, aval_out, x, update,
                                  start_indices=start_indices)
  return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

mlir.register_lowering(dynamic_update_slice_p, _dynamic_update_slice_lower)


def _gather_dtype_rule(operand, indices, *, fill_value, **kwargs):
  if not dtypes.issubdtype(indices.dtype, np.integer):
    raise ValueError("indices must have an integer type")
  return operand.dtype

_rank = lambda arr: len(arr.shape)

def _is_sorted(dims, op_name, name):
  for i in range(1, len(dims)):
    if dims[i] < dims[i - 1]:
      raise TypeError(f"{name} in {op_name} op must be sorted; got {dims}")

def _dims_in_range(dims, rank, op_name, name):
  for dim in dims:
    if dim < 0 or dim >= rank:
      raise TypeError(f"Invalid {name} set in {op_name} op; valid range is "
                      f"[0, {rank}); got: {dim}.")

def _sorted_dims_in_range(dims, rank, op_name, name):
  if len(dims) == 0:
    return
  invalid_dim = None
  if dims[0] < 0:
    invalid_dim = dims[0]
  elif dims[-1] >= rank:
    invalid_dim = dims[-1]
  if invalid_dim:
    raise TypeError(f"Invalid {name} set in {op_name} op; valid range is "
                    f"[0, {rank}); got: {invalid_dim}.")

def _no_duplicate_dims(dims, op_name, name):
  if len(set(dims)) != len(dims):
    raise TypeError(f"{name} in {op_name} op must not repeat; got: {dims}.")

def _disjoint_dims(dims1, dims2, op_name, name1, name2):
  if not set(dims1).isdisjoint(set(dims2)):
    raise TypeError(f"{name1} and {name2} in {op_name} op must be disjoint; "
                    f"got: {dims1} and {dims2}.")

def _gather_shape_rule(operand, indices, *, dimension_numbers,
                       slice_sizes, unique_indices, indices_are_sorted,
                       mode, fill_value):
  """Validates the well-formedness of the arguments to Gather.

  The code implements the checks based on the detailed operation semantics of
  XLA's `Gather <https://www.openxla.org/xla/operation_semantics#gather>`_
  operator and following the outline of the implementation of
  ShapeInference::InferGatherShape in TensorFlow.
  """

  offset_dims = dimension_numbers.offset_dims
  collapsed_slice_dims = dimension_numbers.collapsed_slice_dims
  operand_batching_dims = dimension_numbers.operand_batching_dims
  start_indices_batching_dims = dimension_numbers.start_indices_batching_dims
  start_index_map = dimension_numbers.start_index_map

  # Note: in JAX, index_vector_dim is always computed as below, cf. the
  # documentation of the GatherDimensionNumbers class.
  index_vector_dim = _rank(indices) - 1

  # This case should never happen in JAX, due to the implicit construction of
  # index_vector_dim, but is included for completeness.
  if _rank(indices) < index_vector_dim or index_vector_dim < 0:
    raise TypeError(f"Gather index leaf dimension must be within [0, rank("
                    f"indices) + 1). rank(indices) is {_rank(indices)} and "
                    f"gather index leaf dimension is {index_vector_dim}.")

  # Start ValidateGatherDimensions
  # In the error messages output by XLA, "offset_dims" is called "Output window
  # dimensions" in error messages. For consistency's sake, our error messages
  # stick to "offset_dims".
  _is_sorted(offset_dims, "gather", "offset_dims")
  _no_duplicate_dims(offset_dims, "gather", "offset_dims")

  output_offset_dim_count = len(offset_dims)
  output_shape_rank = len(offset_dims) + _rank(indices) - 1

  for i in range(output_offset_dim_count):
    offset_dim = offset_dims[i]
    if offset_dim < 0 or offset_dim >= output_shape_rank:
      raise TypeError(f"Offset dimension {i} in gather op is out of bounds; "
                      f"got {offset_dim}, but should have been in "
                      f"[0, {output_shape_rank})")

  if len(start_index_map) != indices.shape[index_vector_dim]:
    raise TypeError(f"Gather op has {len(start_index_map)} elements in "
                    f"start_index_map and the bound of dimension "
                    f"{index_vector_dim=} of indices is "
                    f"{indices.shape[index_vector_dim]}. These two "
                    f"numbers must be equal.")

  for i in range(len(start_index_map)):
    operand_dim_for_start_index_i = start_index_map[i]
    if (operand_dim_for_start_index_i < 0 or
        operand_dim_for_start_index_i >= _rank(operand)):
      raise TypeError(f"Invalid start_index_map; domain is "
                      f"[0, {_rank(operand)}), got: "
                      f"{i}->{operand_dim_for_start_index_i}.")

  _no_duplicate_dims(start_index_map, "gather", "start_index_map")

  # _is_sorted and _sorted_dims_in_range are checked in the opposite order
  # compared to the XLA implementation. In cases when the input is not sorted
  # AND there are problematic collapsed_slice_dims, the error message will thus
  # be different.
  _is_sorted(collapsed_slice_dims, "gather", "collapsed_slice_dims")
  _sorted_dims_in_range(collapsed_slice_dims, _rank(operand), "gather",
                        "collapsed_slice_dims")
  _no_duplicate_dims(collapsed_slice_dims, "gather", "collapsed_slice_dims")

  _no_duplicate_dims(operand_batching_dims, "gather", "operand_batching_dims")
  _is_sorted(operand_batching_dims, "gather", "operand_batching_dims")
  _sorted_dims_in_range(
      operand_batching_dims, _rank(operand), "gather", "operand_batching_dims"
  )

  _disjoint_dims(collapsed_slice_dims, operand_batching_dims, "gather",
                 "collapsed_slice_dims", "operand_batching_dims")
  _disjoint_dims(start_index_map, operand_batching_dims, "gather",
                 "start_index_map", "operand_batching_dims")

  _no_duplicate_dims(
      start_indices_batching_dims, "gather", "start_indices_batching_dims"
  )
  _dims_in_range(
      start_indices_batching_dims,
      _rank(indices),
      "gather",
      "start_indices_batching_dims",
  )
  if index_vector_dim in start_indices_batching_dims:
    raise TypeError(
        "Gather op cannot have the index vector dimension as a batching "
        f"dimension; got {start_indices_batching_dims}."
    )

  if len(operand_batching_dims) != len(start_indices_batching_dims):
    raise TypeError(
        "Gather op requires equal numbers of operand_batching_dims and "
        f"start_indices_batching_dims, got {operand_batching_dims} and"
        f"{start_indices_batching_dims}."
    )

  operand_batch_shape = tuple(operand.shape[i] for i in operand_batching_dims)
  indices_batch_shape = tuple(
      indices.shape[i] for i in start_indices_batching_dims
  )
  if not core.definitely_equal_shape(operand_batch_shape, indices_batch_shape):
    raise TypeError(
        "Gather op requires operand batching dimensions and indices batching "
        f"dimensions to have the same shape, got {operand_batch_shape} and "
        f"{indices_batch_shape}."
    )
  # End ValidateGatherDimensions

  if _rank(operand) != len(slice_sizes):
    raise TypeError(f"Gather op must have one slice size for every input "
                    f"dimension; got: len(slice_sizes)={len(slice_sizes)}, "
                    f"input_shape.rank={_rank(operand)}")

  if len(slice_sizes) != len(offset_dims) + len(collapsed_slice_dims) + len(
      operand_batching_dims
  ):
    raise TypeError(
        "All components of the offset index in a gather op must "
        "either be a offset dimension or explicitly collapsed/batching; "
        f"got len(slice_sizes)={len(slice_sizes)}, "
        f"output_slice_sizes={offset_dims}, collapsed_slice_dims="
        f"{collapsed_slice_dims}, operand_batching_dims="
        f"{operand_batching_dims}."
    )

  for i in range(len(slice_sizes)):
    slice_size = slice_sizes[i]
    corresponding_input_size = operand.shape[i]

    if not core.is_empty_shape(indices.shape) and not (
        slice_size >= 0 and corresponding_input_size >= slice_size
    ):
      raise TypeError(f"Slice size at index {i} in gather op is out of range, "
                      f"must be within [0, {corresponding_input_size} + 1), "
                      f"got {slice_size}.")

  for i in range(len(collapsed_slice_dims)):
    bound = slice_sizes[collapsed_slice_dims[i]]
    if bound != 1:
      raise TypeError(f"Gather op can only collapse slice dims with bound 1, "
                      f"but bound is {bound} for index "
                      f"{collapsed_slice_dims[i]} at position {i}.")

  for i in range(len(operand_batching_dims)):
    bound = slice_sizes[operand_batching_dims[i]]
    if bound > 1:
      raise TypeError(f"Gather op can only have operand batching dims with "
                      f"bound 0/1, but bound is {bound} for index "
                      f"{operand_batching_dims[i]} at position {i}."
      )

  return _gather_shape_computation(indices, dimension_numbers, slice_sizes)


def _gather_shape_computation(indices, dimension_numbers, slice_sizes):
  offset_dims = dimension_numbers.offset_dims
  collapsed_slice_dims = dimension_numbers.collapsed_slice_dims
  operand_batching_dims = dimension_numbers.operand_batching_dims
  output_shape_rank = len(offset_dims) + _rank(indices) - 1

  index_vector_dim = _rank(indices) - 1
  expanded_indices_shape = list(indices.shape)

  # This case should never happen in JAX, due to the implicit construction of
  # index_vector_dim, but is included for completeness.
  if len(expanded_indices_shape) == index_vector_dim:
    expanded_indices_shape.append(1)

  expanded_indices_shape.pop(index_vector_dim)

  indices_shape_gen = iter(expanded_indices_shape)

  slice_sizes_gen = (
      s
      for i, s in enumerate(slice_sizes)
      if i not in collapsed_slice_dims and i not in operand_batching_dims
  )
  ans = tuple(next(slice_sizes_gen) if i in offset_dims
              else next(indices_shape_gen) for i in range(output_shape_rank))
  return ans


def _gather_spec_computation(operand, indices, dimension_numbers, slice_sizes):
  """Returns gather output sharding spec if unambiguous, else None.

  Operand dimensions can be split into:
    1. Batching dims which must resolve unambiguously with the corresponding
      indices batching dims. `operand_batching_dims` in GatherDimensionNumbers.
    2. Sliced dims which must be replicated. These are the subset of dims in
      `start_index_map` where the slice size is not equal to the operand size.
      A further subset of these (`collapsed_slice_dims`) are collapsed in the
      output.
    3. Unsliced dims, these correspond directly to dimensions in the
      output and so propagate their shardings. These are the dimensions not
      batched or sliced.

  Indices dimensions can be split into:
    1. Batching dims which must resolve unambiguously with the corresponding
      operand batching dims. `start_indices_batching_dims` in
      GatherDimensionNumbers.
    2. Index vector dim which contains the start indices for each dimension of
      the operand sliced in to. This must be replicated. It is the last
      dimension of indices.
    3. Other dims which correspond directly to dimensions in the output and
      so propagate their shardings. These are the dimensions not present in
      index batching dims and index vector dim.

  If the axes of the corresponding batching dims between operand and indices are
  both not None and do not match, then sharding propagation cannot be resolved
  unambiguously and so we return None.
  """
  offset_dims = dimension_numbers.offset_dims
  start_index_map = dimension_numbers.start_index_map
  collapsed_slice_dims = dimension_numbers.collapsed_slice_dims
  operand_batching_dims = dimension_numbers.operand_batching_dims
  start_indices_batching_dims = dimension_numbers.start_indices_batching_dims
  output_shape_rank = len(offset_dims) + indices.ndim - 1
  index_vector_dim = indices.ndim - 1

  operand_spec = operand.sharding.spec
  indices_spec = list(indices.sharding.spec)

  if (all(s is None for s in operand_spec) and
      all(s is None for s in indices_spec)):
    return P()

  assert all(i in start_index_map for i in collapsed_slice_dims)

  # start_index_map defined which operand dimensions are sliced in to. However,
  # if that slice is a full slice then we can propagate the sharding.
  operand_index_dims_full_slice_or_replicated = all(
      slice_sizes[i] == operand.shape[i] or operand_spec[i] is None
      for i in start_index_map)
  batching_dims_resolve_unambiguously = all(
      indices_spec[indices_dim] == operand_spec[operand_dim]
      or operand_spec[operand_dim] is None
      or indices_spec[indices_dim] is None
      for (operand_dim, indices_dim) in zip(
          operand_batching_dims, start_indices_batching_dims)
  )
  # Leads to comms on the bwd pass in scatter
  indices_non_batch_dims_replicated = all(
      s is None for i, s in enumerate(indices_spec)
      if i not in start_indices_batching_dims
  )

  if (operand_index_dims_full_slice_or_replicated
      and batching_dims_resolve_unambiguously
      and indices_non_batch_dims_replicated):
    # Resolve any batched shardings into indices shardings.
    for operand_dim, indices_dim in zip(
        operand_batching_dims, start_indices_batching_dims):
      assert (indices_spec[indices_dim] == operand_spec[operand_dim]
              or indices_spec[indices_dim] is None
              or operand_spec[operand_dim] is None)
      # Resolution and propagation of batching_dims is handled in indices,
      # operand_batching_dims spec is resolved into indices spec (to match how
      # the gather shape rule resolves output dimensions).
      indices_spec[indices_dim] = indices_spec[indices_dim] or operand_spec[operand_dim]

    slice_sizes_gen = (
        (i, s) for i, s in enumerate(slice_sizes)
        if i not in collapsed_slice_dims and i not in operand_batching_dims)
    indices_spec_gen = iter(indices_spec)

    out_spec = []
    for i in range(output_shape_rank):
      if i in offset_dims:
        # The offset dims are the set of dimensions in the `gather` output that
        # derive solely from the operand.
        operand_dim, slice_size = next(slice_sizes_gen)
        # Due to the `operand_index_dims_full_slice_or_replicated` check, if a
        # slice is not full, it must be replicated.
        assert (slice_size == operand.shape[operand_dim]
                or operand_spec[operand_dim] is None)
        out_spec.append(operand_spec[operand_dim])
      else:
        # The other dimensions are either batching dims (which derive from both
        # indices and operand, and we resolved above) or solely from indices.
        out_spec.append(next(indices_spec_gen))
    return P(*out_spec)
  return None


def _resolve_mesh(*meshes) -> mesh_lib.AbstractMesh:
  """Resolves the mesh between given meshes."""
  unique_meshes = {mesh for mesh in meshes if not mesh.empty}
  if len(unique_meshes) > 1:
    raise core.ShardingTypeError(
        f"Conflicting meshes received. Got: {unique_meshes=}"
    )
  if not unique_meshes:
    return mesh_lib.get_abstract_mesh()
  mesh, = unique_meshes
  return mesh


def _gather_sharding_rule(operand, indices, *, dimension_numbers,
                          slice_sizes, unique_indices, indices_are_sorted,
                          mode, fill_value):
  out_mesh = _resolve_mesh(operand.sharding.mesh, indices.sharding.mesh)
  out_spec = _gather_spec_computation(operand, indices, dimension_numbers,
                                      slice_sizes)
  if out_spec is None:
    raise core.ShardingTypeError(
        "Use `.at[...].get(out_sharding=)` to provide output PartitionSpec for"
        " the gather indexing as out sharding could not be resolved"
        " unambiguously (or would require collectives on inputs). Got"
        f" {operand=}, {indices=}")
  return NamedSharding(out_mesh, out_spec)

def _gather_fill(operand, indices, *, dimension_numbers, slice_sizes,
                 unique_indices, indices_are_sorted, fill_value,
                 output_shape):
  """Lowers a FILL_OR_DROP gather as a PROMISE_IN_BOUNDS gather with masking."""
  dnums = dimension_numbers
  index_dtype = lax_utils.int_dtype_for_shape(operand.shape, signed=True)
  intarray = partial(np.array, dtype=index_dtype)
  operand_dims = lax.shape_as_value(operand.shape).astype(index_dtype)
  indices = lax.convert_element_type(indices, index_dtype)
  num_batch_dims = len(indices.shape) - 1

  upper_bound = operand_dims[
      intarray(dnums.start_index_map)
  ] - lax.shape_as_value(slice_sizes)[intarray(dnums.start_index_map)].astype(
      index_dtype
  )
  mask = lax.bitwise_and(
      lax.ge(indices, index_dtype.type(0)),
      lax.le(indices, lax.expand_dims(upper_bound, tuple(range(num_batch_dims)))))
  mask = lax.reduce_and(mask, [num_batch_dims])

  # Computes the output shape and the positions of the batch dimensions in the
  # output
  output_ndims = num_batch_dims + len(dnums.offset_dims)
  batch_dims_in_output = np.delete(np.arange(output_ndims),
                                   dnums.offset_dims)

  # We don't consume unique_indices directly in gather(), only in its transpose
  # (scatter).
  gather_out = gather(operand, indices, dnums, slice_sizes,
                      indices_are_sorted=indices_are_sorted,
                      mode=GatherScatterMode.PROMISE_IN_BOUNDS)
  return lax.select(
    lax.broadcast_in_dim(mask, output_shape, batch_dims_in_output,
                         out_sharding=gather_out.aval.sharding),
    gather_out, lax.full_like(gather_out, fill_value=fill_value))


def _gather_jvp_rule(g, operand, indices, *, dimension_numbers,
                     slice_sizes, unique_indices, indices_are_sorted, mode,
                     fill_value):
  return gather(g, indices, dimension_numbers, slice_sizes,
                unique_indices=unique_indices,
                indices_are_sorted=indices_are_sorted, mode=mode,
                fill_value=0)

def _gather_transpose_rule(t, operand, indices, *, dimension_numbers,
                           slice_sizes, unique_indices, indices_are_sorted,
                           mode, fill_value):
  assert ad.is_undefined_primal(operand)
  if type(t) is ad_util.Zero:
    out = ad_util.Zero(operand.aval)
  else:
    zeros = lax.full(operand.aval.shape, 0, core.typeof(t).dtype,
                     sharding=operand.aval.sharding)
    zeros = core.pvary(zeros, tuple(operand.aval.vma))
    scatter_dnums = ScatterDimensionNumbers(
        update_window_dims=dimension_numbers.offset_dims,
        inserted_window_dims=dimension_numbers.collapsed_slice_dims,
        scatter_dims_to_operand_dims=dimension_numbers.start_index_map,
        operand_batching_dims=dimension_numbers.operand_batching_dims,
        scatter_indices_batching_dims=dimension_numbers.start_indices_batching_dims,
    )
    out = scatter_add(zeros, indices, t, scatter_dnums,
                      unique_indices=unique_indices,
                      indices_are_sorted=indices_are_sorted,
                      mode=mode)
  return [out, None]

def _gather_batching_rule(batched_args, batch_dims, *, dimension_numbers,
                          slice_sizes, unique_indices, indices_are_sorted,
                          mode, fill_value):
  operand, indices = batched_args
  operand_bdim, indices_bdim = batch_dims

  if operand_bdim is not None and indices_bdim is None:
    operand = batching.moveaxis(operand, operand_bdim, 0)
    operand_bdim = 0
    slice_sizes = (operand.shape[0],) + slice_sizes
    offset_dims = (0,) + tuple(np.add(1, dimension_numbers.offset_dims))
    collapsed_slice_dims = tuple(np.add(1, dimension_numbers.collapsed_slice_dims))
    operand_batching_dims = tuple(
        np.add(1, dimension_numbers.operand_batching_dims)
    )
    start_index_map = tuple(np.add(1, dimension_numbers.start_index_map))
    dnums = GatherDimensionNumbers(
        offset_dims=offset_dims,
        collapsed_slice_dims=collapsed_slice_dims,
        start_index_map=start_index_map,
        operand_batching_dims=operand_batching_dims,
        start_indices_batching_dims=dimension_numbers.start_indices_batching_dims,
    )
    bdim_out = operand_bdim
    return gather(
        operand, indices, dimension_numbers=dnums,
        slice_sizes=slice_sizes,
        unique_indices=unique_indices,
        indices_are_sorted=indices_are_sorted, mode=mode,
        fill_value=fill_value), bdim_out

  elif operand_bdim is None and indices_bdim is not None:
    indices = batching.moveaxis(indices, indices_bdim, 0)
    offset_dims = tuple(1 + d for d in dimension_numbers.offset_dims)
    start_indices_batching_dims = tuple(
        np.add(1, dimension_numbers.start_indices_batching_dims)
    )
    dnums = GatherDimensionNumbers(
        offset_dims=offset_dims,
        collapsed_slice_dims=dimension_numbers.collapsed_slice_dims,
        start_index_map=dimension_numbers.start_index_map,
        operand_batching_dims=dimension_numbers.operand_batching_dims,
        start_indices_batching_dims=start_indices_batching_dims,
    )
    # If batching indexed accesses into the same array, the batched gather may
    # no longer have sorted or unique indices.
    return gather(operand, indices, dimension_numbers=dnums,
                  slice_sizes=slice_sizes, unique_indices=False,
                  indices_are_sorted=False, mode=mode, fill_value=fill_value), 0

  else:
    # move batch dimensions to the front to simplify logic
    operand = batching.moveaxis(operand, operand_bdim, 0)
    indices = batching.moveaxis(indices, indices_bdim, 0)

    if core.definitely_equal(operand.shape[0], 0):
      slice_sizes = (0,) + slice_sizes
    else:
      slice_sizes = (1,) + slice_sizes
    collapsed_slice_dims = tuple(
        np.add(1, dimension_numbers.collapsed_slice_dims)
    )
    operand_batching_dims = (0,) + tuple(
        np.add(1, dimension_numbers.operand_batching_dims)
    )
    start_indices_batching_dims = (0,) + tuple(
        np.add(1, dimension_numbers.start_indices_batching_dims)
    )
    offset_dims = tuple(np.add(1, dimension_numbers.offset_dims))
    start_index_map = tuple(np.add(1, dimension_numbers.start_index_map))

    dnums = GatherDimensionNumbers(
        offset_dims=offset_dims,
        collapsed_slice_dims=collapsed_slice_dims,
        start_index_map=start_index_map,
        operand_batching_dims=operand_batching_dims,
        start_indices_batching_dims=start_indices_batching_dims,
    )
    return gather(operand, indices, dimension_numbers=dnums,
                  slice_sizes=slice_sizes, unique_indices=unique_indices,
                  indices_are_sorted=indices_are_sorted, mode=mode,
                  fill_value=fill_value), 0

gather_p = standard_primitive(
    _gather_shape_rule, _gather_dtype_rule, 'gather',
    weak_type_rule=_argnum_weak_type(0), sharding_rule=_gather_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'gather'))
ad.defjvp(gather_p, _gather_jvp_rule, None)
ad.primitive_transposes[gather_p] = _gather_transpose_rule
batching.primitive_batchers[gather_p] = _gather_batching_rule


def _gather_lower_opaque(ctx, operand, indices, *,
                         dimension_numbers, slice_sizes, unique_indices,
                         indices_are_sorted, mode, fill_value) -> ir.Value:
  aval_x, aval_indices = ctx.avals_in
  aval_y, = ctx.avals_out
  elt_shape = core.physical_element_aval(aval_x.dtype).shape
  trailing_offset_dims = [aval_y.ndim + i for i in range(len(elt_shape))]
  dimension_numbers = dimension_numbers._replace(
      offset_dims=(*dimension_numbers.offset_dims, *trailing_offset_dims))
  slice_sizes = (*slice_sizes, *elt_shape)
  gather_lower = partial(
      _gather_lower, dimension_numbers=dimension_numbers,
      slice_sizes=slice_sizes, unique_indices=unique_indices,
      indices_are_sorted=indices_are_sorted, mode=mode, fill_value=fill_value)
  res, = mlir.delegate_lowering(
      ctx, gather_lower, operand, indices,
      avals_in=[core.physical_aval(aval_x), aval_indices],
      avals_out=[core.physical_aval(aval_y)])
  return res

def _gather_lower(ctx, operand, indices, *,
                  dimension_numbers, slice_sizes, unique_indices,
                  indices_are_sorted, mode, fill_value):
  _, indices_aval = ctx.avals_in
  aval_out, = ctx.avals_out
  if dtypes.issubdtype(aval_out.dtype, dtypes.extended):
    return [_gather_lower_opaque(
        ctx, operand, indices, dimension_numbers=dimension_numbers,
        slice_sizes=slice_sizes, unique_indices=unique_indices,
        indices_are_sorted=indices_are_sorted, mode=mode,
        fill_value=fill_value)]

  if mode == GatherScatterMode.FILL_OR_DROP:
    gather_fill_fn = mlir.lower_fun(_gather_fill, multiple_results=False)
    return gather_fill_fn(
        ctx, operand, indices,
        dimension_numbers=dimension_numbers, slice_sizes=slice_sizes,
        unique_indices=unique_indices, indices_are_sorted=indices_are_sorted,
        fill_value=fill_value, output_shape=aval_out.shape)

  assert mode in (GatherScatterMode.PROMISE_IN_BOUNDS,
                  GatherScatterMode.CLIP), mode
  dnums = hlo.GatherDimensionNumbers.get(
      collapsed_slice_dims=list(dimension_numbers.collapsed_slice_dims),
      operand_batching_dims=list(dimension_numbers.operand_batching_dims),
      start_indices_batching_dims=list(
          dimension_numbers.start_indices_batching_dims
      ),
      index_vector_dim=len(indices_aval.shape) - 1,
      offset_dims=list(dimension_numbers.offset_dims),
      start_index_map=list(dimension_numbers.start_index_map),
  )
  if not core.is_constant_shape(slice_sizes):
    slice_sizes = mlir.eval_dynamic_shape_as_tensor(ctx, slice_sizes)
    # TODO(burmako): Fix overly conservative type inference of DynamicGatherOp.
    # For now use the build_generic so that we can specify the result type.
    # return hlo.DynamicGatherOp(
    #     operand, indices, mlir.shape_tensor(slice_sizes),
    #     dnums, indices_are_sorted=ir.BoolAttr.get(indices_are_sorted)).results
    results = [mlir.aval_to_ir_type(aval_out)]
    operands = [operand, indices, slice_sizes]
    attributes = {
        "dimension_numbers": dnums,
        "indices_are_sorted": ir.BoolAttr.get(indices_are_sorted)
    }
    return hlo.DynamicGatherOp.build_generic(
        results=results, operands=operands, attributes=attributes).results
  elif core.is_empty_shape(aval_out.shape):
    out = mlir.full_like_aval(ctx, 0, aval_out)
    return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]
  else:
    out = hlo.gather(operand, indices, dnums, mlir.dense_int_array(slice_sizes),
                     indices_are_sorted=ir.BoolAttr.get(indices_are_sorted))
    return [mlir.lower_with_sharding_in_types(ctx, out, aval_out)]

mlir.register_lowering(gather_p, _gather_lower)

def _scatter_dtype_rule(operand, indices, updates, **kwargs):
  if not dtypes.issubdtype(indices.dtype, np.integer):
    raise ValueError("indices must have an integer type")
  lax.check_same_dtypes("scatter", operand, updates)
  return operand.dtype

def _get_updates_batching_dims(indices_batching_dims, update_window_dims,
                               index_vector_dim, updates_shape):
  scatter_dim_in_updates = list(range(index_vector_dim))
  for i in update_window_dims:
    scatter_dim_in_updates.insert(i, None)  # type: ignore
  assert len(scatter_dim_in_updates) == len(updates_shape)
  return tuple(scatter_dim_in_updates.index(i) for i in indices_batching_dims)

def _scatter_shape_rule(operand, indices, updates, *, update_jaxpr,
                        update_consts, dimension_numbers, indices_are_sorted,
                        unique_indices, mode):
  """Validates the well-formedness of the ``dimension_numbers`` argument to
  Scatter.

  The code implements the checks based on the detailed operation semantics of
  XLA's `Scatter <https://www.openxla.org/xla/operation_semantics#scatter>`_
  operator and following the outline of the implementation of
  ShapeInference::InferScatterShape in TensorFlow.
  """

  update_window_dims = dimension_numbers.update_window_dims
  inserted_window_dims = dimension_numbers.inserted_window_dims
  operand_batching_dims = dimension_numbers.operand_batching_dims
  scatter_indices_batching_dims = dimension_numbers.scatter_indices_batching_dims
  scatter_dims_to_operand_dims = dimension_numbers.scatter_dims_to_operand_dims
  # Note: in JAX, index_vector_dim is always computed as below, cf. the
  # documentation of the ScatterDimensionNumbers class.
  index_vector_dim = _rank(indices) - 1

  # This case should never happen in JAX, due to the implicit construction of
  # index_vector_dim, but is included for completeness.
  if _rank(indices) < index_vector_dim or index_vector_dim < 0:
    raise TypeError(f"Scatter index leaf dimension must be within [0, "
                    f"rank(indices) + 1). rank(indices) is {_rank(indices)} "
                    f"and scatter index leaf dimension is {index_vector_dim}.")

  expanded_indices_shape = list(indices.shape)
  # This case should never happen in JAX, due to the implicit construction of
  # index_vector_dim, but is included for completeness.
  if len(expanded_indices_shape) == index_vector_dim:
    expanded_indices_shape.append(1)

  expected_updates_rank = (len(expanded_indices_shape) - 1 +
                           len(update_window_dims))

  if _rank(updates) != expected_updates_rank:
    raise TypeError(f"Updates tensor must be of rank {expected_updates_rank}; "
                    f"got {_rank(updates)}.")

  # Validate update_window_dims
  _is_sorted(update_window_dims, "scatter", "update_window_dims")
  _no_duplicate_dims(update_window_dims, "scatter", "update_window_dims")
  _sorted_dims_in_range(update_window_dims, _rank(updates), "scatter",
                        "update_window_dims")

  # Validate inserted_window_dims
  _is_sorted(inserted_window_dims, "scatter", "inserted_window_dims")
  _no_duplicate_dims(inserted_window_dims, "scatter", "inserted_window_dims")
  _sorted_dims_in_range(inserted_window_dims, _rank(operand), "scatter",
                        "inserted_window_dims")

  # Validate operand_batching_dims and scatter_indices_batching_dims
  _is_sorted(operand_batching_dims, "scatter", "operand_batching_dims")
  _no_duplicate_dims(operand_batching_dims, "scatter", "operand_batching_dims")
  _sorted_dims_in_range(
      operand_batching_dims, _rank(operand), "scatter", "operand_batching_dims"
  )
  _disjoint_dims(inserted_window_dims, operand_batching_dims, "scatter",
                 "inserted_window_dims", "operand_batching_dims")
  _disjoint_dims(scatter_dims_to_operand_dims, operand_batching_dims, "scatter",
                 "scatter_dims_to_operand_dims", "operand_batching_dims")

  _no_duplicate_dims(
      scatter_indices_batching_dims, "scatter", "scatter_indices_batching_dims"
  )
  _dims_in_range(
      scatter_indices_batching_dims,
      _rank(indices),
      "scatter",
      "scatter_indices_batching_dims",
  )
  if index_vector_dim in scatter_indices_batching_dims:
    raise TypeError(
        "Scatter op cannot have the index vector dimension as a batching "
        f"dimension; got {scatter_indices_batching_dims}.")

  if len(operand_batching_dims) != len(scatter_indices_batching_dims):
    raise TypeError(
        "Scatter op requires equal numbers of operand_batching_dims and "
        f"scatter_indices_batching_dims, got {operand_batching_dims} and "
        f"{scatter_indices_batching_dims}."
    )

  operand_batch_shape = tuple(operand.shape[i] for i in operand_batching_dims)
  indices_batch_shape = tuple(
      indices.shape[i] for i in scatter_indices_batching_dims
  )
  if not core.definitely_equal_shape(operand_batch_shape, indices_batch_shape):
    raise TypeError(
        "Scatter op requires operand batching dimensions and indices batching "
        f"dimensions to have the same shape, got {operand_batch_shape} and "
        f"{indices_batch_shape}."
    )
  updates_batching_dims = _get_updates_batching_dims(
      scatter_indices_batching_dims, update_window_dims, index_vector_dim,
      updates.shape)
  updates_batch_shape = tuple(updates.shape[i] for i in updates_batching_dims)
  if not core.definitely_equal_shape(operand_batch_shape, updates_batch_shape):
    raise TypeError(
        "Scatter op requires operand batching dimensions and updates batching "
        f"dimensions to have the same shape, got {operand_batch_shape} and "
        f"{updates_batch_shape}."
    )
  # Validate window_size
  window_size = (
      len(update_window_dims) +
      len(inserted_window_dims) +
      len(operand_batching_dims)
  )
  if _rank(operand) != window_size:
    raise TypeError(f"Scatter op has window of size {window_size}; doesn't "
                    f"match operand of rank {_rank(operand)}.")

  # Validate scatter_dims_to_operand_dims
  if (len(scatter_dims_to_operand_dims) !=
      indices.shape[index_vector_dim]):
    raise TypeError(f"Scatter op has {len(scatter_dims_to_operand_dims)} "
                    f"elements in scatter_dims_to_operand_dims and the bound "
                    f"of dimension {index_vector_dim=} of "
                    f"indices is {indices.shape[index_vector_dim]}. These two "
                    f"numbers must be equal")

  for i in range(len(scatter_dims_to_operand_dims)):
    dim = scatter_dims_to_operand_dims[i]
    if dim < 0 or dim >= _rank(operand):
      raise TypeError(f"Invalid scatter_dims_to_operand_dims mapping; domain "
                      f"is [0, {_rank(operand)}), got: {i}->{dim}.")

  _no_duplicate_dims(scatter_dims_to_operand_dims, "scatter",
                     "scatter_dims_to_operand_dims")

  max_update_slice_sizes = [
      operand.shape[i]
      for i in range(len(operand.shape))
      if (
          i not in set(inserted_window_dims)
          and i not in set(operand_batching_dims)
      )
  ]

  for i in range(len(update_window_dims)):
    update_window_dim = update_window_dims[i]
    if max_update_slice_sizes[i] < updates.shape[update_window_dim]:
      raise TypeError(f"Bounds of the window dimensions of updates must not "
                      f"exceed the bounds of the corresponding dimensions of "
                      f"operand. For dimension {update_window_dim}, updates "
                      f"bound is {updates.shape[update_window_dim]}, operand "
                      f"bound is {max_update_slice_sizes[i]}.")

  update_scatter_dims = [dim for dim in range(_rank(updates)) if dim not in
                         set(update_window_dims)]

  scatter_dims_seen = 0
  for i in update_scatter_dims:
    if scatter_dims_seen == index_vector_dim:
      scatter_dims_seen += 1
    if not core.definitely_equal(updates.shape[i], expanded_indices_shape[scatter_dims_seen]):
      raise TypeError(f"Bounds of the scatter dimensions of updates must be "
                      f"the same as the bounds of the corresponding dimensions "
                      f"of scatter indices. For scatter dimension {i}, updates "
                      f"bound is {updates.shape[i]}, indices bound is "
                      f"{expanded_indices_shape[scatter_dims_seen]}.")
    scatter_dims_seen += 1

  return operand.shape


def _is_resolvable(*axis_names: str | None) -> bool:
  """Checks if given sharding axis names resolve unambiguously."""
  assert len(axis_names) >= 2, "At least two axis names expected."
  return len({a for a in axis_names if a is not None}) <= 1


def _scatter_spec_computation(
    operand, indices, updates, dimension_numbers) -> P | None:
  """For a scatter, we consider the gather rules in inverse.

  We consider a gather, then convert to scatter as the inverse.
  A gather is a group of queries, each query slices a window out of `operand`.

  The `operand` dims sliced in to are those in `start_index_map`. Some of size 1
  dims in `start_index_map` are squeezed out of the gather output, termed
  `collapsed_slice_dims`. Hence `collapsed_slice_dims` is a subset of
  `start_index_map`. Confusingly, some of the slice sizes into `operand` can be
  full slices, hence it can be inferred that the corresponding start index of
  the window in that dim is 0 and hence the `operand` is not sliced in to.

  All dims of `indices`, except the final dim (`index_vector_dim`) are different
  queries in to `operand`. These queries may be batched with `operand`, and so
  are effectively N (query, operand) pairs so the queries are non-overlapping,
  each into different matching sized slices of `operand` (N being the batch
  dimension size), or unbatched where N queries are into the same `operand`. The
  actual indices of the start of the (fixed size) windows are contained in the
  final dimension of `indices`, the `index_vector_dim`.

  Some dims of `operand` are not sliced in to, these are `offset_dims`
  [offset_dims is defined from the gather output perspective (the `updates` in
  scatter) and by construction are the frontmost dims in `operand` after
  disregarding `collapsed_slice_dims` and `operand_batching_dims`].

  For a gather:
   - Batch dimensions must be resolvable unambiguously between `operand` and
   `indices`
   - Dims the `operand` is sliced in to are either full slices, or the `operand`
    is replicated in that dim (all the data for each query is present)
   - The `index_vector_dim` containing the actual start indices of the windows
   in to `operand` must be replicated

  A scatter is a group of queries, each of which has a corresponding window of
  `updates` to put in to `operand`.

  For a scatter:
   1 - Batch dimensions must resolve unambiguously between `operand`, `indices`
   and `updates`
   2 - Full slice dims must resolve unambiguously between `operand` and
   `updates`
   3 - Sub slice dims in `operand` and `updates` must be replicated.
   4 - `Indices` and `updates` must be replicated in dims where both are
   updating, possibly overlapping, subslices of operand. These are referred to
   as unbatched queries in the above description of gather.
   5 - The `index_vector_dim` containing the actual start indices of the windows
   in to `operand` must be replicated

  Not all dimensions updating slices in `operand` exist in `updates`. The
  `inserted_window_dims` provide size 1 updates into `operand` and are inserted
  into `updates`. We handle `operand` slices present in `updates`, then handle
  the inserted dims separately.

  Correspondance:
                Gather <-> Scatter
           offset_dims <-> update_window_dims
  collapsed_slice_dims <-> inserted_window_dims
       start_index_map <-> scatter_dims_to_operand_dims
  """
  operand_batching_dims = dimension_numbers.operand_batching_dims
  indices_batching_dims = dimension_numbers.scatter_indices_batching_dims
  update_window_dims = dimension_numbers.update_window_dims
  inserted_window_dims = dimension_numbers.inserted_window_dims
  index_vector_dim = indices.ndim - 1

  operand_spec = operand.sharding.spec
  indices_spec = indices.sharding.spec
  updates_spec = updates.sharding.spec

  if (all(s is None for s in operand_spec) and
      all(s is None for s in indices_spec) and
      all(s is None for s in updates_spec)):
    return P()

  updates_batching_dims = _get_updates_batching_dims(
      indices_batching_dims, update_window_dims, index_vector_dim, updates.shape)

  # Work out the corresponding operand dim for update window dims
  operand_window_dims = [
      i for i in range(operand.ndim)
      if i not in inserted_window_dims and i not in operand_batching_dims]

  # 1 - Batch dimensions must resolve unambiguously between `operand`, `indices`
  # and `updates`
  batch_dims_resolvable = all(
      _is_resolvable(operand_spec[od], indices_spec[id], updates_spec[ud])
      for od, id, ud in zip(
          operand_batching_dims, indices_batching_dims, updates_batching_dims))

  # 2 - Full slice dims must resolve unambiguously between `operand` and `updates`
  # 3 - Sub slice dims in `operand` and `updates` must be replicated.
  update_and_operand_window_dims_resolvable = all(
      (updates.shape[update_dim] == operand.shape[operand_dim] and
       updates_spec[update_dim] == operand_spec[operand_dim])
      or (updates_spec[update_dim] is None and operand_spec[operand_dim] is None)
      for update_dim, operand_dim in zip(
          update_window_dims, operand_window_dims))
  inserted_window_dims_replicated_in_operand = all(
      spec is None for i, spec in enumerate(operand_spec)
      if i in inserted_window_dims)

  # 4 - `Indices` and `updates` must be replicated in dims where both are
  # updating, possibly overlapping, slices of operand. These are referred to as
  # unbatched queries in the above description of gather.
  unbatched_query_dims_in_updates_replicated = all(
      spec is None for i, spec in enumerate(updates_spec)
      if i not in updates_batching_dims and i not in update_window_dims)
  unbatched_query_dims_in_indices_replicated = all(
      spec is None for i, spec in enumerate(indices_spec[:index_vector_dim])
      if i not in indices_batching_dims)

  # 5 - The `index_vector_dim` containing the actual start indices of the
  # windows in to `operand` must be replicated
  index_vector_dim_is_replicated = indices_spec[index_vector_dim] is None

  if (batch_dims_resolvable and
      update_and_operand_window_dims_resolvable and
      inserted_window_dims_replicated_in_operand and
      unbatched_query_dims_in_updates_replicated and
      unbatched_query_dims_in_indices_replicated and
      index_vector_dim_is_replicated):
    out_spec = list(operand_spec)

    # 1 - Batch dims
    for operand_dim, indices_dim, updates_dim in zip(
        operand_batching_dims, indices_batching_dims, updates_batching_dims):
      if out_spec[operand_dim] is None:
        out_spec[operand_dim] = (
            indices_spec[indices_dim] or updates_spec[updates_dim])
    # 2, 3 - Full/sub slices of operand dims present in updates
    for update_dim, operand_dim in zip(update_window_dims, operand_window_dims):
      if out_spec[operand_dim] is None:
        out_spec[operand_dim] = updates_spec[update_dim]
    return P(*out_spec)

  return None


def _scatter_memory_space_rule(
    operand, indices, updates, *, update_jaxpr, update_consts,
    dimension_numbers, indices_are_sorted, unique_indices, mode):
  return operand.memory_space


def _scatter_sharding_rule(
    operand, indices, updates, *, update_jaxpr, update_consts,
    dimension_numbers, indices_are_sorted, unique_indices, mode):
  out_mesh = _resolve_mesh(
      *(x.sharding.mesh for x in (operand, indices, updates)))
  out_spec = _scatter_spec_computation(operand, indices, updates,
                                       dimension_numbers)
  if out_spec is None:
    raise core.ShardingTypeError(
        "Use `.at[...].set/add/mul/...(out_sharding=)` to provide output"
        " PartitionSpec for the scatter update as out sharding could not be"
        " resolved unambiguously (or would require collectives on inputs). Got"
        f" {operand=}, {indices=}, {updates=}")
  return NamedSharding(out_mesh, out_spec)

def _clamp_scatter_indices(operand, indices, updates, *, dnums):
  """Clamps `indices` to be in-range for a scatter."""
  slice_sizes = []
  pos = 0
  for i in range(len(operand.shape)):
    if i in dnums.inserted_window_dims or i in dnums.operand_batching_dims:
      slice_sizes.append(1)
    else:
      slice_sizes.append(updates.shape[dnums.update_window_dims[pos]])
      pos += 1

  upper_bounds: core.Shape = tuple(operand.shape[i] - slice_sizes[i]
                                   for i in dnums.scatter_dims_to_operand_dims)

  # Stack upper_bounds into a Array[n]
  upper_bound = lax.shape_as_value(upper_bounds)
  # This fix fails lax_test_no_jax_array
  upper_bound = lax.min(
      upper_bound,
      upper_bound.dtype.type(
          min(np.iinfo(upper_bound.dtype).max, np.iinfo(indices.dtype).max)
      ),
  )
  upper_bound = lax.convert_element_type(upper_bound, indices.dtype)
  upper_bound = lax.broadcast_in_dim(upper_bound, indices.shape,
                                     (len(indices.shape) - 1,))
  return lax.clamp(indices.dtype.type(0), indices, upper_bound)

def _scatter_addsub_jvp(
    prim, primals, tangents, *, update_jaxpr, update_consts, dimension_numbers,
    indices_are_sorted, unique_indices, mode):
  operand, indices, updates = primals
  g_operand, g_indices, g_updates = tangents
  del g_indices  # ignored
  val_out = prim.bind(
      operand, indices, updates, update_jaxpr=update_jaxpr,
      update_consts=update_consts, dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=mode)
  if type(g_operand) is ad_util.Zero and type(g_updates) is ad_util.Zero:
    tangent_out = ad_util.Zero.from_primal_value(val_out)
  else:
    g_operand = ad.instantiate_zeros(g_operand)
    g_updates = ad.instantiate_zeros(g_updates)
    tangent_out = prim.bind(
        g_operand, indices, g_updates, update_jaxpr=update_jaxpr,
        update_consts=update_consts, dimension_numbers=dimension_numbers,
        indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
        mode=mode)
  return val_out, tangent_out


def _scatter_addsub_transpose_rule(
    prim, t, operand, indices, updates, *, update_jaxpr, update_consts,
    dimension_numbers, indices_are_sorted, unique_indices, mode):
  assert not ad.is_undefined_primal(indices)
  if ad.is_undefined_primal(updates):
    updates_shape = updates.aval.shape
  else:
    updates_shape = updates.shape
  if type(t) is ad_util.Zero:
    operand_t = ad_util.Zero(operand.aval) if ad.is_undefined_primal(operand) else None
    update_t = ad_util.Zero(updates.aval) if ad.is_undefined_primal(updates) else None
  else:
    operand_t = update_t = None
    if ad.is_undefined_primal(operand):
      operand_t = t

    if ad.is_undefined_primal(updates):
      gather_dnums = GatherDimensionNumbers(
          offset_dims=dimension_numbers.update_window_dims,
          collapsed_slice_dims=dimension_numbers.inserted_window_dims,
          start_index_map=dimension_numbers.scatter_dims_to_operand_dims,
          operand_batching_dims=dimension_numbers.operand_batching_dims,
          start_indices_batching_dims=dimension_numbers.scatter_indices_batching_dims,
      )
      slice_sizes = []
      pos = 0
      for i in range(len(t.shape)):
        if (
            i in dimension_numbers.inserted_window_dims
            or i in dimension_numbers.operand_batching_dims
        ):
          slice_sizes.append(1)
        else:
          slice_sizes.append(updates_shape[dimension_numbers.update_window_dims[pos]])
          pos += 1
      update_t = gather(t, indices, dimension_numbers=gather_dnums,
                        slice_sizes=slice_sizes, mode=mode, fill_value=0)
      if prim is scatter_sub_p:
        update_t = lax.neg(update_t)
  return [operand_t, None, update_t]

def _scatter_mul_transpose_rule(t, operand, indices, updates, *,
                                update_jaxpr, update_consts, dimension_numbers,
                                indices_are_sorted, unique_indices, mode):
  assert not ad.is_undefined_primal(indices)
  if ad.is_undefined_primal(updates):
    updates_shape = updates.aval.shape
  else:
    updates_shape = updates.shape
  if type(t) is ad_util.Zero:
    operand_t = ad_util.Zero(operand.aval) if ad.is_undefined_primal(operand) else None
    update_t = ad_util.Zero(updates.aval) if ad.is_undefined_primal(updates) else None
  else:
    operand_t = update_t = None
    if ad.is_undefined_primal(operand):
      operand_t = scatter_mul(
          t, indices, updates, dimension_numbers=dimension_numbers,
          indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
          mode=mode)
    if ad.is_undefined_primal(updates):
      if not unique_indices:
        raise NotImplementedError(
          "scatter_mul gradients are only implemented if `unique_indices=True`")
      gather_dnums = GatherDimensionNumbers(
          offset_dims=dimension_numbers.update_window_dims,
          collapsed_slice_dims=dimension_numbers.inserted_window_dims,
          start_index_map=dimension_numbers.scatter_dims_to_operand_dims,
          operand_batching_dims=dimension_numbers.operand_batching_dims,
          start_indices_batching_dims=dimension_numbers.scatter_indices_batching_dims,
      )
      slice_sizes = []
      pos = 0
      for i in range(len(t.shape)):
        if (
            i in dimension_numbers.inserted_window_dims
            or i in dimension_numbers.operand_batching_dims
        ):
          slice_sizes.append(1)
        else:
          slice_sizes.append(updates_shape[dimension_numbers.update_window_dims[pos]])
          pos += 1
      update_t = gather(lax.mul(t, operand), indices,
                        dimension_numbers=gather_dnums, slice_sizes=slice_sizes,
                        mode=mode, fill_value=0)
  return [operand_t, None, update_t]


def _scatter_batching_rule(scatter_op, axis_data, batched_args, batch_dims, *,
                           update_jaxpr, update_consts, dimension_numbers,
                           indices_are_sorted, unique_indices, mode):
  operand, indices, updates = batched_args
  operand_bdim, indices_bdim, updates_bdim = batch_dims

  # move the operand batch dim to the front if it is not None, otherwise create
  # it at the front (so that we can scatter into it)
  size = next(x.shape[ax] for x, ax in zip(batched_args, batch_dims)
              if ax is not None)
  operand = batching.bdim_at_front(operand, operand_bdim, size,
                                   axis_data.explicit_mesh_axis)
  updates = batching.bdim_at_front(updates, updates_bdim, size,
                                   axis_data.explicit_mesh_axis)

  if indices_bdim is None:
    inserted_window_dims = tuple(np.add(1, dimension_numbers.inserted_window_dims))
    update_window_dims = (0,) + tuple(np.add(1, dimension_numbers.update_window_dims))
    operand_batching_dims = tuple(
        np.add(1, dimension_numbers.operand_batching_dims)
    )
    scatter_dims_to_operand_dims = tuple(np.add(1, dimension_numbers.scatter_dims_to_operand_dims))
    dnums = ScatterDimensionNumbers(
        update_window_dims=update_window_dims,
        inserted_window_dims=inserted_window_dims,
        scatter_dims_to_operand_dims=scatter_dims_to_operand_dims,
        operand_batching_dims=operand_batching_dims,
        scatter_indices_batching_dims=dimension_numbers.scatter_indices_batching_dims,
    )
    return scatter_op.bind(
      operand, indices, updates, dimension_numbers=dnums,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=mode, update_jaxpr=update_jaxpr, update_consts=update_consts), 0

  # see the third case in _gather_batching_rule for comparison and comments
  indices = batching.bdim_at_front(indices, indices_bdim, size,
                                   axis_data.explicit_mesh_axis)

  update_window_dims = tuple(np.add(1, dimension_numbers.update_window_dims))
  inserted_window_dims = tuple(np.add(1, dimension_numbers.inserted_window_dims))
  operand_batching_dims = (0,) + tuple(
      np.add(1, dimension_numbers.operand_batching_dims))
  scatter_indices_batching_dims = (0,) + tuple(
      np.add(1, dimension_numbers.scatter_indices_batching_dims))
  scatter_dims_to_operand_dims = tuple(
      np.add(1, dimension_numbers.scatter_dims_to_operand_dims))
  dnums = ScatterDimensionNumbers(
      update_window_dims=update_window_dims,
      inserted_window_dims=inserted_window_dims,
      scatter_dims_to_operand_dims=scatter_dims_to_operand_dims,
      operand_batching_dims=operand_batching_dims,
      scatter_indices_batching_dims=scatter_indices_batching_dims)
  return scatter_op.bind(
      operand, indices, updates, dimension_numbers=dnums,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=mode, update_jaxpr=update_jaxpr, update_consts=update_consts), 0

scatter_add_p = standard_primitive(
    _scatter_shape_rule, _scatter_dtype_rule, 'scatter-add',
    weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'scatter_add'),
    memory_space_rule=_scatter_memory_space_rule)
ad.primitive_jvps[scatter_add_p] = partial(_scatter_addsub_jvp, scatter_add_p)
ad.primitive_transposes[scatter_add_p] = partial(_scatter_addsub_transpose_rule, scatter_add_p)
batching.fancy_primitive_batchers[scatter_add_p] = partial(_scatter_batching_rule, scatter_add_p)
batching.skippable_batchers[scatter_add_p] = lambda _: ()

scatter_sub_p = standard_primitive(
    _scatter_shape_rule, _scatter_dtype_rule, 'scatter-sub',
    weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'scatter_sub'),
    memory_space_rule=_scatter_memory_space_rule
)
ad.primitive_jvps[scatter_sub_p] = partial(_scatter_addsub_jvp, scatter_sub_p)
ad.primitive_transposes[scatter_sub_p] = partial(_scatter_addsub_transpose_rule, scatter_sub_p)
batching.fancy_primitive_batchers[scatter_sub_p] = partial(
    _scatter_batching_rule, scatter_sub_p)
batching.skippable_batchers[scatter_sub_p] = lambda _: ()

scatter_mul_p = standard_primitive(
    _scatter_shape_rule, _scatter_dtype_rule, 'scatter-mul',
    weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'scatter_mul'),
    memory_space_rule=_scatter_memory_space_rule)

def _scatter_mul_jvp_rhs(g, x, i, y, *, dimension_numbers,
                         indices_are_sorted, unique_indices, mode, **kw):
  if not unique_indices:
    raise NotImplementedError(
      "scatter_mul gradients are only implemented if `unique_indices=True`")
  return lax.mul(x, scatter_add(
      ad_util.zeros_like_jaxval(x), i, g, dimension_numbers=dimension_numbers,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=mode))

ad.defjvp(scatter_mul_p,
          lambda g, x, i, y, **kw: scatter_mul_p.bind(g, i, y, **kw),
          None,
          _scatter_mul_jvp_rhs)
ad.primitive_transposes[scatter_mul_p] = _scatter_mul_transpose_rule
batching.fancy_primitive_batchers[scatter_mul_p] = (
  partial(_scatter_batching_rule, scatter_mul_p))
batching.skippable_batchers[scatter_mul_p] = lambda _: ()

def _scatter_extremal_jvp(scatter_op, primals, tangents, update_jaxpr,
                          update_consts, dimension_numbers,
                          indices_are_sorted, unique_indices, mode):
  operand, indices, updates = primals
  g_operand, g_indices, g_updates = tangents

  scatter_dnums = dimension_numbers
  updates_shape = updates.shape

  val_out = scatter_op.bind(
      operand, indices, updates, update_jaxpr=update_jaxpr,
      update_consts=update_consts, dimension_numbers=scatter_dnums,
      indices_are_sorted=indices_are_sorted,
      unique_indices=unique_indices, mode=mode)

  if type(g_operand) is ad_util.Zero and type(g_updates) is ad_util.Zero:
    tangent_out = ad_util.Zero.from_primal_value(val_out)
  else:
    g_operand = ad.instantiate_zeros(g_operand)
    g_updates = ad.instantiate_zeros(g_updates)

    # gather_dnums and slice_sizes define the gather op that is the inverse of
    # the scatter op specified by scatter_dnums
    gather_dnums = GatherDimensionNumbers(
        offset_dims=scatter_dnums.update_window_dims,
        collapsed_slice_dims=scatter_dnums.inserted_window_dims,
        start_index_map=scatter_dnums.scatter_dims_to_operand_dims,
        operand_batching_dims=scatter_dnums.operand_batching_dims,
        start_indices_batching_dims=scatter_dnums.scatter_indices_batching_dims,
    )

    slice_sizes = []
    pos = 0
    for i in range(len(operand.shape)):
      if (
          i in scatter_dnums.inserted_window_dims
          or i in scatter_dnums.operand_batching_dims
      ):
        slice_sizes.append(1)
      else:
        slice_sizes.append(updates_shape[scatter_dnums.update_window_dims[pos]])
        pos += 1

    # For consistency with other max operations, if there are two or more values
    # in updates that are contending to replace the same index location, the
    # resulting tangent at that location will be the average of the associated
    # tangents for the values in updates.

    initial_vals = gather(
        operand, indices, gather_dnums, np.array(slice_sizes))

    target_vals = gather(
        val_out, indices, gather_dnums, np.array(slice_sizes))

    successful_updates = (updates == target_vals)
    retained_values = (initial_vals == target_vals)

    num_updates = gather(
        scatter_add(
            lax._zeros(operand), indices,
            lax.select(successful_updates, lax._ones(updates),
                       lax._zeros(updates)),
            scatter_dnums),
        indices,
        gather_dnums,
        np.array(slice_sizes))

    num_refs = gather(
        scatter_add(lax._zeros(operand),
                    indices,
                    lax._ones(updates),
                    scatter_dnums),
        indices,
        gather_dnums,
        np.array(slice_sizes))

    updates_normalizer = lax.select(retained_values,
                                    1.0 / (num_updates + 1),
                                    1.0 / num_updates)

    updates_coef = lax.select(successful_updates,
                              updates_normalizer,
                              lax._zeros(updates))

    operand_normalizer = lax.select(retained_values,
                                    1.0 / (num_updates + 1),
                                    lax._zeros(num_updates))

    operand_coef = (-1.0 + operand_normalizer) / num_refs

    # This can be simplified once scatter has transpose implemented
    target_tangents = gather(
        g_operand, indices, gather_dnums, np.array(slice_sizes))

    tangent_updates = (target_tangents * operand_coef +
                       g_updates * updates_coef)

    tangent_out = scatter_add(g_operand,
                              indices,
                              tangent_updates,
                              scatter_dnums,
                              indices_are_sorted=indices_are_sorted,
                              unique_indices=unique_indices,
                              mode=mode)

  return val_out, tangent_out

scatter_min_p = standard_primitive(
    _scatter_shape_rule, _scatter_dtype_rule, 'scatter-min',
    weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'scatter_min'),
    memory_space_rule=_scatter_memory_space_rule)
batching.fancy_primitive_batchers[scatter_min_p] = (
  partial(_scatter_batching_rule, scatter_min_p))
batching.skippable_batchers[scatter_min_p] = lambda _: ()
ad.primitive_jvps[scatter_min_p] = partial(_scatter_extremal_jvp, scatter_min_p)

scatter_max_p = standard_primitive(
    _scatter_shape_rule, _scatter_dtype_rule, 'scatter-max',
    weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'scatter_max'),
    memory_space_rule=_scatter_memory_space_rule)
batching.fancy_primitive_batchers[scatter_max_p] = (
  partial(_scatter_batching_rule, scatter_max_p))
batching.skippable_batchers[scatter_max_p] = lambda _: ()
ad.primitive_jvps[scatter_max_p] = partial(_scatter_extremal_jvp, scatter_max_p)

def _scatter_jvp(primals, tangents, *, update_jaxpr, update_consts,
                 dimension_numbers, indices_are_sorted, unique_indices,
                 mode):
  if update_jaxpr is not None:
    raise NotImplementedError("scatter_apply JVP not implemented")
  operand, indices, updates = primals
  g_operand, g_indices, g_updates = tangents
  dnums = dimension_numbers

  if type(g_operand) is ad_util.Zero and type(g_updates) is ad_util.Zero:
    val_out = scatter_p.bind(
      operand, indices, updates, update_jaxpr=update_jaxpr,
      update_consts=update_consts, dimension_numbers=dnums,
      indices_are_sorted=indices_are_sorted, unique_indices=unique_indices,
      mode=mode)
    return val_out, ad_util.Zero.from_primal_value(val_out)

  g_operand = ad.instantiate_zeros(g_operand)
  g_updates = ad.instantiate_zeros(g_updates)

  if unique_indices:
    # If the user has promised that the updates don't overlap, we can use a much
    # simpler JVP.
    val_out = scatter_p.bind(
      operand, indices, updates, update_jaxpr=update_jaxpr,
      update_consts=update_consts, dimension_numbers=dnums,
      indices_are_sorted=indices_are_sorted, unique_indices=True, mode=mode)
    tangent_out = scatter_p.bind(
      g_operand, indices, g_updates, update_jaxpr=update_jaxpr,
      update_consts=update_consts, dimension_numbers=dnums,
      indices_are_sorted=indices_are_sorted, unique_indices=True, mode=mode)
    return val_out, tangent_out

  # If there are overlapping indices in the scatter, it is unspecified which
  # update "wins". So we use the following perhaps surprising scheme:
  # a) attach a positive ID to each update in updates, and perform the scatter
  #    on the IDs
  # b) perform the inverse gather on the scattered IDs (similar to
  #    _scatter_add_transpose).
  # c) use the gathered IDs to mask the primal and tangent values.
  # d) perform a scatter-add on the masked primal and tangent values. A benefit
  #    of using scatter-add here is that we don't need a `scatter` transpose
  #    rule.

  # a) attach a positive ID to each update in `updates`, and perform a scatter
  #    on the IDs.
  ids_shape = list(updates.shape)
  for update_dim in dnums.update_window_dims:
    ids_shape[update_dim] = 1
  num_ids = math.prod(ids_shape)
  id_dtype = lax_utils.int_dtype_for_dim(num_ids, signed=False)
  update_ids = lax.add(lax.reshape(lax.iota(id_dtype, num_ids), ids_shape),
                       lax._ones(updates, dtype=id_dtype))

  scattered_ids = scatter(lax.full(operand.shape, 0, id_dtype),
                          indices, update_ids, dnums,
                          indices_are_sorted=indices_are_sorted,
                          unique_indices=unique_indices, mode=mode)

  # b) compute the inverse gather that "undoes" the scatter on the id values.
  gather_dnums = GatherDimensionNumbers(
      offset_dims=dnums.update_window_dims,
      collapsed_slice_dims=dnums.inserted_window_dims,
      start_index_map=dnums.scatter_dims_to_operand_dims,
      operand_batching_dims=dnums.operand_batching_dims,
      start_indices_batching_dims=dnums.scatter_indices_batching_dims,
  )
  slice_sizes = []
  pos = 0
  for i in range(len(scattered_ids.shape)):
    if i in dnums.inserted_window_dims or i in dnums.operand_batching_dims:
      slice_sizes.append(1)
    else:
      slice_sizes.append(updates.shape[dnums.update_window_dims[pos]])
      pos += 1
  gathered_update_ids = gather(scattered_ids, indices,
                               dimension_numbers=gather_dnums,
                               slice_sizes=slice_sizes)

  # c) mask off input elements that do not correspond to a primal output.
  masked_operand = lax.select(lax.eq(scattered_ids, lax._zeros(scattered_ids)),
                              operand, lax._zeros(operand))
  masked_updates = lax.select(lax.eq(update_ids,  gathered_update_ids),
                              updates, lax._zeros(updates))
  masked_g_operand = lax.select(lax.eq(scattered_ids, lax._zeros(scattered_ids)),
                                g_operand, lax._zeros(g_operand))
  masked_g_updates = lax.select(lax.eq(update_ids, gathered_update_ids),
                                g_updates, lax._zeros(g_updates))

  # d) perform scatter-adds to compute the primal and tangent outputs.
  val_out = scatter_add(masked_operand, indices, masked_updates,
                        dimension_numbers=dnums,
                        indices_are_sorted=indices_are_sorted,
                        unique_indices=unique_indices, mode=mode)
  tangent_out = scatter_add(masked_g_operand, indices, masked_g_updates,
                            dimension_numbers=dnums,
                            indices_are_sorted=indices_are_sorted,
                            unique_indices=unique_indices, mode=mode)
  return val_out, tangent_out

def _scatter_transpose_rule(t, operand, indices, updates, *,
                            update_jaxpr, update_consts, dimension_numbers,
                            indices_are_sorted, unique_indices, mode):
  if not unique_indices:
    raise NotImplementedError("scatter transpose is only implemented where "
                              "unique_indices=True")
  assert not ad.is_undefined_primal(indices)
  if ad.is_undefined_primal(updates):
    updates_shape = updates.aval.shape
    updates_sharding = updates.aval.sharding
  else:
    updates_shape = updates.shape
    updates_sharding = core.typeof(updates).sharding
  if type(t) is ad_util.Zero:
    operand_t = ad_util.Zero(operand.aval) if ad.is_undefined_primal(operand) else None
    update_t = ad_util.Zero(updates.aval) if ad.is_undefined_primal(updates) else None
  else:
    operand_t = update_t = None
    if ad.is_undefined_primal(operand):
      # Zero out gradient entries that correspond to updated indices.
      operand_t = scatter(
          t, indices,
          lax.full(updates_shape, 0, dtype=t.dtype, sharding=updates_sharding),
          dimension_numbers=dimension_numbers,
          indices_are_sorted=indices_are_sorted, unique_indices=True, mode=mode)

    if ad.is_undefined_primal(updates):
      gather_dnums = GatherDimensionNumbers(
          offset_dims=dimension_numbers.update_window_dims,
          collapsed_slice_dims=dimension_numbers.inserted_window_dims,
          start_index_map=dimension_numbers.scatter_dims_to_operand_dims,
          operand_batching_dims=dimension_numbers.operand_batching_dims,
          start_indices_batching_dims=dimension_numbers.scatter_indices_batching_dims,
      )
      slice_sizes = []
      pos = 0
      for i in range(len(t.shape)):
        if (
            i in dimension_numbers.inserted_window_dims
            or i in dimension_numbers.operand_batching_dims
        ):
          slice_sizes.append(1)
        else:
          slice_sizes.append(updates_shape[dimension_numbers.update_window_dims[pos]])
          pos += 1
      update_t = gather(t, indices, dimension_numbers=gather_dnums,
                        slice_sizes=slice_sizes, mode=mode,
                        fill_value=0)

  return [operand_t, None, update_t]

scatter_p = standard_primitive(
    _scatter_shape_rule, _scatter_dtype_rule, 'scatter',
    weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
    vma_rule=partial(core.standard_vma_rule, 'scatter'),
    memory_space_rule=_scatter_memory_space_rule)
ad.primitive_jvps[scatter_p] = _scatter_jvp
ad.primitive_transposes[scatter_p] = _scatter_transpose_rule
batching.fancy_primitive_batchers[scatter_p] = (
  partial(_scatter_batching_rule, scatter_p))
batching.skippable_batchers[scatter_p] = lambda _: ()


def _scatter_lower_opaque(ctx, operand, indices, updates, *,
                          update_jaxpr, update_consts, dimension_numbers,
                          unique_indices, indices_are_sorted, mode):
  aval_x, aval_indices, aval_updates = ctx.avals_in
  aval_y, = ctx.avals_out
  elt_shape = core.physical_element_aval(aval_x.dtype).shape
  trailing_window_dims = [aval_updates.ndim + i for i in range(len(elt_shape))]
  dimension_numbers = dimension_numbers._replace(
      update_window_dims=(*dimension_numbers.update_window_dims,
                          *trailing_window_dims))
  scatter_lower = partial(
      _scatter_lower, update_jaxpr=update_jaxpr, update_consts=update_consts,
      dimension_numbers=dimension_numbers, unique_indices=unique_indices,
      indices_are_sorted=indices_are_sorted, mode=mode)
  res, = mlir.delegate_lowering(
      ctx, scatter_lower, operand, indices, updates,
      avals_in=[core.physical_aval(aval_x), aval_indices,
                core.physical_aval(aval_updates)],
      avals_out=[core.physical_aval(aval_y)])
  return res


def _scatter_lower(ctx: mlir.LoweringRuleContext, operand, indices, updates, *,
                   update_jaxpr: core.Jaxpr, update_consts, dimension_numbers,
                   indices_are_sorted, unique_indices, mode):
  if update_jaxpr is None:
    assert not update_consts
    operand_dtype = ctx.avals_in[0].dtype
    update_jaxpr, update_consts = lax._reduction_jaxpr(
        _scatter_reduction_computation, core.ShapedArray((), operand_dtype))

  aval_out, = ctx.avals_out
  if dtypes.issubdtype(aval_out.dtype, dtypes.extended):
    return [_scatter_lower_opaque(
        ctx, operand, indices, updates,
        update_jaxpr=update_jaxpr, update_consts=update_consts,
        dimension_numbers=dimension_numbers, unique_indices=unique_indices,
        indices_are_sorted=indices_are_sorted, mode=mode)]

  if mode == GatherScatterMode.CLIP:
    clip_fn = mlir.lower_fun(_clamp_scatter_indices, multiple_results=False)
    (indices,) = clip_fn(ctx.replace(avals_out=None), operand, indices,
                          updates, dnums=dimension_numbers)

  dnums = dimension_numbers
  scatter_dnums = hlo.ScatterDimensionNumbers.get(
      update_window_dims=list(dnums.update_window_dims),
      inserted_window_dims=list(dnums.inserted_window_dims),
      input_batching_dims=list(dnums.operand_batching_dims),
      scatter_indices_batching_dims=list(dnums.scatter_indices_batching_dims),
      scattered_dims_to_operand_dims=list(dnums.scatter_dims_to_operand_dims),
      index_vector_dim=len(ctx.avals_in[1].shape) - 1,
  )
  result = mlir.aval_to_ir_type(aval_out)
  operand = [operand]
  updates = [updates]
  op = hlo.ScatterOp((result,), operand, indices, updates, scatter_dnums,
                     indices_are_sorted=ir.BoolAttr.get(indices_are_sorted),
                     unique_indices=ir.BoolAttr.get(unique_indices))
  scalar_type = mlir.aval_to_ir_type(core.ShapedArray((), aval_out.dtype))
  update = op.update_computation.blocks.append(scalar_type, scalar_type)
  with ir.InsertionPoint(update):
    name_stack = source_info_util.new_name_stack()
    if update_jaxpr.effects:
      raise NotImplementedError('Cannot lower effectful `scatter`.')
    out_nodes, _ = mlir.jaxpr_subcomp(
        ctx.module_context, update_jaxpr, name_stack, mlir.TokenSet(),
        update_consts, update.arguments[0], update.arguments[1],
        dim_var_values=ctx.dim_var_values, const_lowering=ctx.const_lowering)
    hlo.return_(mlir.flatten_ir_values(out_nodes))
  return [mlir.lower_with_sharding_in_types(ctx, r, aval)
          for r, aval in safe_zip(op.results, ctx.avals_out)]

mlir.register_lowering(scatter_p, _scatter_lower)
mlir.register_lowering(scatter_add_p, _scatter_lower)
mlir.register_lowering(scatter_sub_p, _scatter_lower)
mlir.register_lowering(scatter_mul_p, _scatter_lower)
mlir.register_lowering(scatter_min_p, _scatter_lower)
mlir.register_lowering(scatter_max_p, _scatter_lower)


def _real_dtype(dtype): return np.finfo(dtype).dtype


def _scatter_addsub_lower_gpu(
    ctx, operand, indices, updates, *, update_jaxpr, update_consts,
    dimension_numbers, indices_are_sorted, unique_indices, mode,
    reduce_op):
  operand_aval_in, _, updates_aval_in = ctx.avals_in
  if operand_aval_in.dtype != np.complex128:
    return _scatter_lower(ctx, operand, indices, updates,
                          update_jaxpr=update_jaxpr,
                          update_consts=update_consts,
                          dimension_numbers=dimension_numbers,
                          indices_are_sorted=indices_are_sorted,
                          unique_indices=unique_indices, mode=mode)

  if mode == GatherScatterMode.CLIP:
    clip_fn = mlir.lower_fun(_clamp_scatter_indices, multiple_results=False)
    indices, = clip_fn(ctx.replace(avals_out=None), operand, indices, updates,
                       dnums=dimension_numbers)

  aval_out, = ctx.avals_out
  dnums = dimension_numbers
  scatter_dnums = hlo.ScatterDimensionNumbers.get(
      update_window_dims=list(dnums.update_window_dims),
      inserted_window_dims=list(dnums.inserted_window_dims),
      input_batching_dims=list(dnums.operand_batching_dims),
      scatter_indices_batching_dims=list(dnums.scatter_indices_batching_dims),
      scattered_dims_to_operand_dims=list(dnums.scatter_dims_to_operand_dims),
      index_vector_dim=len(ctx.avals_in[1].shape) - 1,
  )
  real_dtype = _real_dtype(aval_out.dtype)
  operand_type_part = mlir.aval_to_ir_type(
      core.ShapedArray(aval_out.shape, real_dtype))

  def _scatter(operand_part, updates_part):
    operand_part = [operand_part]
    updates_part = [updates_part]

    scatter = hlo.ScatterOp(
        (operand_type_part,), operand_part, indices, updates_part, scatter_dnums,
        indices_are_sorted=ir.BoolAttr.get(indices_are_sorted),
        unique_indices=ir.BoolAttr.get(unique_indices))
    scalar_type = mlir.aval_to_ir_type(core.ShapedArray((), real_dtype))
    reducer = scatter.regions[0].blocks.append(scalar_type, scalar_type)
    with ir.InsertionPoint(reducer):
      hlo.return_([reduce_op(*reducer.arguments).result])
    return mlir.lower_with_sharding_in_types(ctx, scatter.result, aval_out)

  real = _scatter(hlo.real(operand), hlo.real(updates))
  imag = _scatter(hlo.imag(operand), hlo.imag(updates))
  return [hlo.complex(real, imag)]


mlir.register_lowering(
    scatter_add_p,
    partial(_scatter_addsub_lower_gpu, reduce_op=hlo.AddOp),
    platform="gpu",
)
mlir.register_lowering(
    scatter_sub_p,
    partial(_scatter_addsub_lower_gpu, reduce_op=hlo.SubtractOp),
    platform="gpu",
)


def _dynamic_slice_indices(
    operand: Array | np.ndarray,
    start_indices: Array | np.ndarray | Sequence[ArrayLike],
    allow_negative_indices: bool | Sequence[bool],
  ) -> list[ArrayLike]:
  # Normalize the start_indices w.r.t. operand.shape
  if len(start_indices) != operand.ndim:
    msg = ("Length of slice indices must match number of operand dimensions ({} "
          "vs {})")
    raise ValueError(msg.format(len(start_indices), operand.shape))
  if not isinstance(start_indices, (tuple, list)):
    if start_indices.ndim != 1:  # type: ignore[union-attr]
      raise ValueError("Slice indices must be a 1D sequence, got {}"
                       .format(start_indices.shape))  # type: ignore[union-attr]
    start_indices = list(start_indices)
  result: list[ArrayLike] = []
  if isinstance(allow_negative_indices, bool):
    allow_negative_indices = [allow_negative_indices] * len(start_indices)
  # Loop to correct for negative indices.
  for i, d, allow_negative_index in zip(
      start_indices, operand.shape, allow_negative_indices
  ):
    # If i is unsigned, then it cannot be negative.
    if dtypes.issubdtype(_dtype(i), np.unsignedinteger):
      result.append(i)
      continue
    # Test whether i and d are static to avoid unnecessary staging.
    if isinstance(i, (int, np.integer)) and core.is_constant_dim(d):
      if allow_negative_index:
        result.append(lax.convert_element_type(i + d if i < 0 else i, _dtype(i)))
      elif i < 0:
        raise ValueError(f"Index {i} is out of bounds for dimension {d} if "
                          "allow_negative_indices=False")
      else:
        result.append(lax.convert_element_type(i, _dtype(i)))
      continue
    d = core.dimension_as_value(d)
    if isinstance(i, (int, np.integer)):
      if allow_negative_index:
        result.append(i + lax.convert_element_type(d, _dtype(i)) if i < 0 else i)
      elif i < 0:
        raise ValueError(f"Index {i} is out of bounds for dimension {d} if "
                          "allow_negative_indices=False")
      else:
        result.append(i)
      continue
    if allow_negative_index:
      d_arr = lax.convert_element_type(d, _dtype(i))
      result.append(lax.select(i < 0, i + d_arr, i))
    else:
      result.append(i)
  return result
