# Copyright 2020 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.
"""Workarounds for jax2tf transforms when XLA is not linked in."""
import builtins
import dataclasses
from functools import partial, wraps
import math
import string
from typing import Any, Callable, Optional, Sequence

from jax._src import core
from jax import lax
from jax._src.lax import slicing as lax_slicing
from jax._src import dtypes
from jax._src import util

from jax.experimental.jax2tf import jax2tf

import numpy as np
import tensorflow as tf  # type: ignore[import]


# Implementation rules for primitives when XLA is not linked in. These
# implementations are workarounds, making use of TF ops that do work when XLA is
# not linked in. They are only used when the argument `enable_xla=False` when
# calling jax2tf.convert().
tf_impl_no_xla: dict[core.Primitive, Callable[..., Any]] = {}


TfVal = Any
DType = Any
PrecisionType = Any


def _error(primitive_name: str, suffix_msg: str = "") -> Exception:
  msg = f"Call to {primitive_name} cannot be converted with enable_xla=False."
  if suffix_msg:
    msg += (f" {suffix_msg} - See source code for the precise conditions under "
             "which it can be converted without XLA.")
  return NotImplementedError(msg)

_conv_error = lambda msg: _error("conv_general_dilated", msg)
_reduce_error = lambda msg: _error("reduce_window", msg)
_scatter_error = lambda msg: _error("scatter_(update/add/multiply/min/max)", msg
                                   )

def _unimplemented(name):

  def op(*arg, **kwargs):
    raise _error(name)

  return op


# TODO(marcvanzee): Remove this function and use `tf.math.invert_permutation`
# once it is implemented by TFjs:
# https://github.com/tensorflow/tfjs/issues/6395.
def _invert_permutation(perm):
  return tuple(perm.index(i) for i in range(len(perm)))


def _transpose_with_shape(x: TfVal, x_shape: core.Shape, permutation) -> tuple[TfVal, core.Shape]:
  """Computes transposition of x and its shape.

  x_shape matches x.shape in the known dimensions, and it has dimension
  polynomials elsewhere, while x.shape has None.
  """
  return tf.transpose(x, perm=permutation), tuple(x_shape[i] for i in permutation)


def _transpose_for_tf_conv(lhs, lhs_shape: core.Shape,
                           rhs, rhs_shape: core.Shape, dimension_numbers):
  """Tranposes lhs and rhs to respectively NHWC and HWIO so they can be passed to TF functions.

  The shapes passed in and returned may contain polynomials, and thus may
  be different than lhs.shape and rhs.shape.
  """
  # TODO(marcvanzee): Add tests for this ops for shape polymorphism.
  lhs_perm, rhs_perm, _ = dimension_numbers

  # TODO(marcvanzee): Consider merging tranposes if we want to optimize.
  # For `lhs_perm` / `output_perm`, perm (0, 1, 2, 3) corresponds to "NCHW".
  lhs, lhs_shape = _transpose_with_shape(lhs, lhs_shape, lhs_perm)  # lhs --> "NCHW"
  if len(lhs_perm) == 3:
    # For 1D convolution, we add a trivial "W" dimension, so that 2D Convolution
    # logic can be applied downstream.
    lhs = lhs[:, :, :, np.newaxis]
    lhs_shape = tuple(lhs_shape) + (1,)
  # However, the TF ops only support "NHWC" on CPU, so we transpose again.
  lhs, lhs_shape = _transpose_with_shape(lhs, lhs_shape, (0, 2, 3, 1))  # "NCHW" --> "NHWC"

  # For `rhs_perm`, perm (0, 1, 2, 3) corresponds to "OIHW".
  rhs, rhs_shape = _transpose_with_shape(rhs, rhs_shape, rhs_perm)  # rhs --> "OIHW"
  # Handle conv1d case.
  if len(rhs_perm) == 3:
    rhs = rhs[:, :, :, np.newaxis]
    rhs_shape = tuple(rhs_shape) + (1,)
  # For the tf ops, rhs is expected to be "OIHW".
  rhs, rhs_shape = _transpose_with_shape(rhs, rhs_shape, (2, 3, 1, 0))  # "OIHW" --> "HWIO"
  jax2tf._assert_matching_abstract_shape(lhs, lhs_shape)
  jax2tf._assert_matching_abstract_shape(rhs, rhs_shape)
  return lhs, lhs_shape, rhs, rhs_shape


def pads_to_padtype(in_shape, window_shape, window_strides, padding) -> str:
  for pad_str in ["VALID", "SAME"]:
    pads = lax.padtype_to_pads(in_shape, window_shape, window_strides, pad_str)
    if list(pads) == list(padding):
      return pad_str
  return "EXPLICIT"


def _pad_spatial_dims(x, x_shape, padding):
  """Pads `x` using `padding`, which specifies padding for the spatial dimensions."""
  padding = tuple(padding)
  if len(padding) == len(x_shape) - 2:
    # If necessary, add empty padding for batch and feature dimensions.
    no_pad = ((0, 0),)
    padding = no_pad + padding + no_pad
  x = tf.pad(x, padding)
  assert len(x.shape) == len(padding)
  x_shape = tuple(p0 + xs + p1 for xs, (p0, p1) in zip(x_shape, padding))
  jax2tf._assert_matching_abstract_shape(x, x_shape)
  return x, x_shape

def _check_pad_spatial_dims(x, x_shape, padding):
  """Pads `x` using `padding`, which specifies padding for the spatial dimensions."""
  padding = tuple(padding)
  if len(padding) == len(x_shape) - 2:
    # If necessary, add empty padding for batch and feature dimensions.
    no_pad = ((0, 0),)
    padding = no_pad + padding + no_pad
  assert len(x.shape) == len(padding)
  x_shape = tuple(p0 + xs + p1 for xs, (p0, p1) in zip(x_shape, padding))
  return x, x_shape, padding

def _conv_transpose_pads_to_padtype(kernel_sdims, lhs_dilation, padding):
  """Finds the padding type for a transpose convolution."""
  # This is simply checking agreement with lax._conv_transpose_padding.
  is_valid = True
  is_same = True
  if not len(kernel_sdims) == len(lhs_dilation) == len(padding):
    raise ValueError(f'Found different lengths for '
                     f'kernel_sdims ({kernel_sdims}), '
                     f'lhs_dilation ({lhs_dilation}), '
                     f'and padding ({padding}).')
  for k, s, (begin, end) in zip(kernel_sdims, lhs_dilation, padding):
    # Check for VALID padding.
    pad_len_valid = k + s - 2 + builtins.max(k - s, 0)
    pad_a = k - 1
    pad_b = pad_len_valid - pad_a
    if begin != pad_a or end != pad_b:
      is_valid = False

    # Check for SAME padding.
    pad_len_same = k + s - 2
    if s > k - 1:
      pad_a = k - 1
    else:
      pad_a = int(np.ceil(pad_len_same / 2))
    pad_b = pad_len_same - pad_a
    if begin != pad_a or end != pad_b:
      is_same = False

  if is_valid:
    return 'VALID'
  elif is_same:
    return 'SAME'
  raise ValueError('Transpose convolution padding mode must be '
                   '`SAME` or `VALID`.')

def _validate_spatial_dimensions(lhs: TfVal, lhs_shape: core.Shape,
                                 rhs: TfVal, rhs_shape: core.Shape):
  """Check spatial dimension support."""
  jax2tf._assert_matching_abstract_shape(lhs, lhs_shape)
  jax2tf._assert_matching_abstract_shape(rhs, rhs_shape)

  nr_spatial_dimensions = len(lhs_shape) - 2
  # Currently we only support 1D+2D convolutions because it keeps the code
  # relatively simple and covers most cases.
  if nr_spatial_dimensions > 2:
    raise _conv_error(
        "We only support 1D or 2D convolutions, but found "
        f"{nr_spatial_dimensions}.")


def _normalize_padding_and_dilations(
      padding, lhs_dilation, rhs_dilation, is_conv1d):
  if is_conv1d:
    lhs_dilation = list(lhs_dilation) + [1]
    rhs_dilation = list(rhs_dilation) + [1]
    # Empty padding in the dummy dimension.
    # Note that when kernel_size=stride=1, padding of (0, 0) is both 'VALID' and
    # 'SAME'. So the inferred padding type will still register according to the
    # first dimension padding.
    padding = list(padding) + [(0, 0)]
  return padding, lhs_dilation, rhs_dilation


def _normalize_window_strides(window_strides):
  """Ensure window_strides has length 4."""
  # Some TF ops require len(window_strides) == 4 while others do not. We simply
  # ensure it always has len(4).
  if len(window_strides) == 1:
    # This is the Conv1D case. We add a dummy dimension to allow using 2D ops,
    # and use stride=1 on the dummy dimension.
    window_strides = list(window_strides) + [1]
  if len(window_strides) == 2:
    window_strides = [1] + list(window_strides) + [1]
  return window_strides


def _validate_conv_features(
      is_transpose, is_atrous, is_depthwise, feature_group_count,
      batch_group_count, preferred_element_type, lhs_dtype):
  if feature_group_count > 1 and not is_depthwise:
    raise _conv_error("Grouped convolutions are unsupported")
  if (is_depthwise and is_atrous) and not is_transpose:
    # We allow dilated depthwise convolutions.
    pass
  elif [is_depthwise, is_atrous, is_transpose].count(True) > 1:
    raise _conv_error(
        f"Can only do one of depthwise ({is_depthwise}), atrous ({is_atrous}) "
        f"and tranposed convolutions ({is_transpose})")

  # We can implement batch grouping when there is a need for it.
  if batch_group_count != 1:
    raise _conv_error("Unimplemented support for batch_group_count != 1 "
                f"(found {batch_group_count})")

  if (preferred_element_type is not None and
      preferred_element_type != lhs_dtype):
    raise _conv_error("Unimplemented support for preferred_element_type")


def _conv_general_dilated(
    lhs, rhs, *, window_strides, padding, lhs_dilation, rhs_dilation,
    dimension_numbers: lax.ConvDimensionNumbers, feature_group_count: int,
    batch_group_count: int,
    precision: Optional[tuple[PrecisionType, PrecisionType]],
    preferred_element_type: Optional[DType],
    _in_avals: Sequence[core.ShapedArray], _out_aval: core.ShapedArray):
  """Implementation of lax.conv_general_dilated_p using XlaConv."""
  # In presence of shape polymorphism, lhs.shape and rhs.shape may contain
  # None. The actual dimension polynomial shapes are in _in_avals.
  del precision  # Unused arguments.
  lhs_shape, rhs_shape = _in_avals[0].shape, _in_avals[1].shape
  out_shape = _out_aval.shape
  _validate_spatial_dimensions(lhs, lhs_shape, rhs, rhs_shape)
  is_conv1d = len(lhs_shape) - 2 == 1

  tf_window_strides = _normalize_window_strides(window_strides)
  padding, lhs_dilation, rhs_dilation = _normalize_padding_and_dilations(
      padding, lhs_dilation, rhs_dilation, is_conv1d)

  lhs, lhs_shape, rhs, rhs_shape = _transpose_for_tf_conv(lhs, lhs_shape,
                                                          rhs, rhs_shape,
                                                          dimension_numbers)
  in_channels = lhs_shape[-1]
  *rhs_spatial_shapes, _, rhs_out_channel = rhs_shape

  is_transpose = any([d != 1 for d in lhs_dilation])
  is_atrous = any([d != 1 for d in rhs_dilation])
  is_depthwise = in_channels == feature_group_count and feature_group_count > 1
  _validate_conv_features(is_transpose, is_atrous, is_depthwise,
                          feature_group_count, batch_group_count,
                          preferred_element_type, lhs.dtype.as_numpy_dtype)

  rhs_dilated_shape = [
      (k - 1) * r + 1 for k, r in zip(rhs_spatial_shapes, rhs_dilation)
  ]
  output_perm = dimension_numbers[2]

  if is_transpose:
    padding_type = _conv_transpose_pads_to_padtype(
        rhs_spatial_shapes, lhs_dilation, padding)
  else:
    padding_type = pads_to_padtype(
      lhs_shape[1:3], rhs_dilated_shape, window_strides, padding)
    # We only manually pad if we aren't using a tranposed convolutions.
    if padding_type == "EXPLICIT":
      lhs, lhs_shape, padding = _check_pad_spatial_dims(lhs, lhs_shape, padding)
      padding_type = padding

  if padding_type != "SAME" and any(l < r for l, r in zip(lhs_shape[1:3], rhs_dilated_shape)):
    # If the input shape is smaller than the filter shape in a spatial dimension,
    # lax returns only zeros while tf.conv2d returns an error.
    # We thus return zeros to make sure the behavior is consistent.
    return tf.broadcast_to(tf.constant(0, dtype=tf.float32),
                           jax2tf._eval_shape(out_shape))

  if is_depthwise:
    # Reshape filter from
    # [filter_height, filter_width, 1, in_channels * channel_multiplier] to
    # [filter_height, filter_width, in_channels, channel_multiplier].
    new_rhs_shape = tuple(rhs_spatial_shapes) + (in_channels,
                                                 rhs_out_channel // in_channels)
    output = tf.nn.depthwise_conv2d(
        input=lhs,
        filter=tf.reshape(rhs, jax2tf._eval_shape(new_rhs_shape)),
        strides=tf_window_strides,
        padding=padding_type,
        dilations=rhs_dilation)

  elif is_transpose:
    # tf.nn.conv2d_transpose requires a transposed filter.
    rhs_t = tf.reverse(rhs, [0, 1])
    rhs_t = tf.transpose(rhs_t, (0, 1, 3, 2))

    # We should transpose `out_shape` to "NHWC", which is what TF expects.
    # First transpose to "NCHW".
    if is_conv1d:
      tf_out_shape = tuple(out_shape[i] for i in output_perm) + (1,)
    else:
      tf_out_shape = tuple(out_shape[i] for i in output_perm)
    # Then transpose "NCHW" to "NHWC".
    tf_out_shape = tuple(tf_out_shape[i] for i in (0, 2, 3, 1))
    output = tf.nn.conv2d_transpose(
        input=lhs,
        filters=rhs_t,
        output_shape=jax2tf._eval_shape(tf_out_shape),
        strides=lhs_dilation,
        padding=padding_type)

  else:
    output = tf.nn.conv2d(
        input=lhs,
        filters=rhs,
        strides=tf_window_strides,
        padding=padding_type,
        dilations=rhs_dilation)

  # TF outputs in format "NHWC", so convert to "NCHW", which is lax's default
  # format.
  output = tf.transpose(output, (0, 3, 1, 2))  # "NHWC" --> "NCHW"
  if is_conv1d:
    output = output[:, :, :, 0]
  # To determine the right permutation, we compute the inverse permutation of
  # `output_perm`, so that when `output_perm` is applied to `output`, we obtain
  # the outpt in NCHW format.
  inverse_perm = _invert_permutation(output_perm)
  output = tf.transpose(output, inverse_perm)  # "NCHW" -> desired output shape.
  return output


tf_impl_no_xla[lax.conv_general_dilated_p] = _conv_general_dilated


def _dot_general(lhs, rhs, *, dimension_numbers,
                 precision: Optional[tuple[PrecisionType, PrecisionType]],
                 preferred_element_type: Optional[DType],
                 _in_avals: Sequence[core.ShapedArray],
                 _out_aval: core.ShapedArray):
  """Implementation of lax.dot_general_p in terms of tf.linalg.einsum."""
  # Unused arguments.
  del precision
  del preferred_element_type

  (lhs_contracting, rhs_contracting), (lhs_batch, rhs_batch) = dimension_numbers
  lhs_ndim, rhs_ndim = len(lhs.shape), len(rhs.shape)

  # This condition ensures that:
  # 1) the batch dimensions are ordered in the same way in lhs and rhs (this is
  #    not strictly necessary, but we would have to reshape the array if that
  #    were not the case;
  # 2) lhs and rhs have the same number of dimensions +/- 1
  # 3) the number of non-batch dimensions in both tensors is either 1 or 2
  # 4) the contracting dimensions are consistent with those of a classic
  #    matrix/matrix, vector/matrix or matrix/vector multiplication.
  if (lhs_batch == rhs_batch == tuple(range(len(lhs_batch))) and
      lhs_ndim - rhs_ndim in [-1, 0, 1] and
      1 <= lhs_ndim - len(lhs_batch) <= 2 and
      1 <= rhs_ndim - len(rhs_batch) <= 2 and
      lhs_contracting == (len(lhs.shape) - 1,) and
      rhs_contracting == (len(lhs_batch),)):
    # All the inputs to tf.linalg.matmul must have 2 inner dimensions,
    # after their batch dimensions, so we need to expand the dimensions
    # appropriately. We can get to this branch with three combinations of
    # inner shapes:
    # - lhs.inner_shape == [a, b], rhs.inner_shape == [b, c]
    #   - in this case, the resulting inner shape is [a, c];
    # - lhs.inner_shape == [b]   , rhs.inner_shape == [b, c]
    #   - in this case, we need to expand lhs to [1, b], and the resulting
    #     shape is [c]. We need to squeeze the result of tf.linalg.matmul
    #     as it will have shape [1, c];
    # - lhs.shape == [batch] + [a, b], rhs.shape == [batch] + [b]
    #   - in this case, we need to expand rhs to [b, 1], and the resulting
    #     shape is [a]. We need to squeeze the result of tf.linalg.matmul
    #     as it will have shape [a, 1];
    # - lhs.shape == [batch] + [b]   , rhs.shape == [batch] + [b]
    #   - in this case, we need to expand lhs to [1, b] and rhs to [b, 1],
    #     and the resulting shape is (). We need to squeeze the result of
    #     tf.linalg.matmul as it will have shape [1, 1].
    squeeze_idxs = []
    if lhs_ndim - len(lhs_batch) == 1:
      lhs = tf.expand_dims(lhs, lhs_ndim - 1)
      squeeze_idxs.append(len(lhs.shape) - 2)
    if rhs_ndim - len(rhs_batch) == 1:
      rhs = tf.expand_dims(rhs, rhs_ndim)
      squeeze_idxs.append(len(rhs.shape) - 1)
    result = tf.linalg.matmul(lhs, rhs)
    if len(squeeze_idxs) != 0:
      assert all([result.shape[i] == 1 for i in squeeze_idxs])
      result = tf.squeeze(result, squeeze_idxs)
    return result

  new_id = iter(string.ascii_letters)
  lhs_axis_ids = [next(new_id) for _ in lhs.shape]
  rhs_axis_ids = [next(new_id) for _ in rhs.shape]
  lhs_out_axis_ids = lhs_axis_ids[:]
  rhs_out_axis_ids = rhs_axis_ids[:]

  for lhs_axis, rhs_axis in zip(lhs_contracting, rhs_contracting):
    shared_id = next(new_id)
    lhs_axis_ids[lhs_axis] = shared_id
    rhs_axis_ids[rhs_axis] = shared_id
    lhs_out_axis_ids[lhs_axis] = None  # type: ignore[call-overload]
    rhs_out_axis_ids[rhs_axis] = None  # type: ignore[call-overload]

  batch_ids = []
  for lhs_axis, rhs_axis in zip(lhs_batch, rhs_batch):
    shared_id = next(new_id)
    lhs_axis_ids[lhs_axis] = shared_id
    rhs_axis_ids[rhs_axis] = shared_id
    lhs_out_axis_ids[lhs_axis] = None  # type: ignore[call-overload]
    rhs_out_axis_ids[rhs_axis] = None  # type: ignore[call-overload]
    batch_ids.append(shared_id)

  not_none = lambda x: x is not None
  out_axis_ids = list(
      filter(not_none, batch_ids + lhs_out_axis_ids + rhs_out_axis_ids))
  assert lhs.dtype == rhs.dtype
  spec = "{},{}->{}".format("".join(lhs_axis_ids), "".join(rhs_axis_ids),
                            "".join(out_axis_ids))
  return tf.linalg.einsum(spec, lhs, rhs)


tf_impl_no_xla[lax.dot_general_p] = _dot_general


def _interior_padding(operand, padding_value, padding_config, operand_shape):
  # Used only when enable_xla=False
  # Applies only the interior padding from the padding_config.
  # We do this somewhat inefficiently, as as a scatter.
  # For each dimension we compute the indices_by_dim as [0, f, 2f, 3f, ...] where
  # f is the dilation factor for the dimension, i.e., 1 + interior_padding.
  # Then we compute the cartesian production of the indices (using broadcast
  # and concat).

  # We could make this code more complex and do all the padding at once, but
  # we prefer to keep it simple.
  indices_by_dim = []
  indices_shape = operand_shape + (1,)
  output_shape = []  # considering only interior padding
  for d, (dsz, (_, _, i)) in enumerate(zip(operand_shape, padding_config)):
    dilation_factor = i + 1
    output_shape.append(dsz * dilation_factor - i)
    indices = tf.range(dsz) * dilation_factor
    expansion = [None] * (1 + len(operand_shape))
    expansion[d] = slice(None, None, None)
    indices_by_dim.append(tf.broadcast_to(indices[expansion], indices_shape))

  indices_cartesian = tf.concat(indices_by_dim, axis=len(operand_shape))
  scattered = tf.scatter_nd(indices_cartesian, operand, output_shape)
  # What elements from the output array we use from
  mask = tf.scatter_nd(indices_cartesian, tf.ones_like(operand, dtype=np.bool_),
                       output_shape)
  return tf.where(mask, scattered, padding_value)


def _pad(operand, padding_value, *, padding_config,
         _in_avals: Sequence[core.ShapedArray], _out_aval: core.ShapedArray):
  # Do only the interior padding first. This is rarely needed.
  if any(i != 0 for _, _, i in padding_config):
    operand = _interior_padding(operand, padding_value, padding_config,
                                jax2tf._eval_shape(_in_avals[0].shape))

  # Now do the non-negative edge padding. This is the common case, use tf.pad.
  non_negative_padding = [((lo if lo >= 0 else 0), (hi if hi >= 0 else 0))
                          for lo, hi, _ in padding_config]
  operand = tf.pad(
      operand,
      non_negative_padding,
      mode="CONSTANT",
      constant_values=padding_value)
  # Now the negative edge padding (this is also rare)
  if any(lo < 0 or hi < 0 for lo, hi, _ in padding_config):
    output_shape = jax2tf._eval_shape(_out_aval.shape)
    begins = [(-lo if lo < 0 else 0) for lo, _, _ in padding_config]
    operand = tf.slice(operand, begins, output_shape)

  return operand


tf_impl_no_xla[lax.pad_p] = _pad


def _argminmax(is_min: bool, operand: TfVal, axes: Sequence[int],
               index_dtype: DType, _in_avals: Sequence[core.ShapedArray],
               _out_aval: core.ShapedArray):
  # The following is known to diverge from JAX behavior for NaN.
  axis, = axes
  output_type = tf.int32
  if dtypes.iinfo(index_dtype).bits > 32:
    output_type = tf.int64
  # TODO(phawkins): handle axes larger than 2^31.
  fn = tf.math.argmin if is_min else tf.math.argmax
  result = fn(operand, axis=axis, output_type=output_type)
  return tf.cast(result, jax2tf._to_tf_dtype(index_dtype))


tf_impl_no_xla[lax.argmin_p] = partial(_argminmax, True)
tf_impl_no_xla[lax.argmax_p] = partial(_argminmax, False)


def _validate_reduce_window_inputs(operand_shape, computation_name, dtype,
                                   window_dimensions, window_strides,
                                   base_dilation, window_dilation):
  if computation_name not in ["min", "max", "add"]:
    raise _reduce_error("Reduction function should be either min, max, or add.")
  if computation_name in ["min", "max"] and dtype in [
      tf.bool, tf.uint32, tf.uint64, tf.complex64, tf.complex128
  ]:
    raise _reduce_error("Min/max pool does not support operands of type "
                        f"{dtype}")
  if computation_name == "min" and dtype in [tf.uint8, tf.uint16]:
    # TODO(marcvanzee): We currently implement min pooling by negating the
    # input, but this doesn't work for uint. We could work around it using
    # tf.math.reduce_min.
    raise _reduce_error(f"Min pool does not support operands of type {dtype}")
  if computation_name == "add" and dtype not in [
      tf.float16, tf.float32, tf.float64
  ]:
    raise _reduce_error("Add pooling does not support operands of type "
                        f"{dtype}")

  if (len(operand_shape) != len(window_dimensions) != len(window_strides) !=
      len(window_dilation)):
    raise _reduce_error("Input shapes, window dimensions, window stride "
                        "dimensions, and window dilation dimensions should "
                        "match.")

  has_only_spatial_dims = True
  if len(operand_shape) > 4:
    raise _reduce_error("Only 1D or 2D input are supported.")
  if len(operand_shape) > 2:
    # operand_shape = (batch, spatial_dims, ..., channel).
    has_only_spatial_dims = False

    for name, value in [("window_dimensions", window_dimensions),
                        ("window_strides", window_strides),
                        ("window_dilation", window_dilation)]:
      if value[0] != value[-1] != 1:
        raise _reduce_error("Only 1D or 2D input are supported, expected "
                            f"{name}=(1, spatial_dims, ..., 1), but got "
                            f"{value}")

  if list(base_dilation) != [1] * len(operand_shape):
    # TODO(marcvanzee): Add support for base dilations. We can do this using
    # a scatter on operand.
    raise _reduce_error("Unimplemented support for base dilation.")

  return has_only_spatial_dims


def _padding_reduce_window(operand, operand_shape, computation_name,
                           window_dimensions, window_strides, padding):
  padding_type = pads_to_padtype(operand_shape, window_dimensions,
                                 window_strides, padding)

  # https://github.com/google/jax/issues/11874.
  needs_manual_padding = (
      padding_type == "SAME" and computation_name == "add" and
      window_dimensions != [1] * len(operand_shape))

  if needs_manual_padding or padding_type == "EXPLICIT":
    operand, operand_shape = _pad_spatial_dims(operand, operand_shape, padding)
    padding_type = "VALID"

  return operand, operand_shape, padding_type


def _reshape_reduce_window(operand, operand_shape, window_dimensions,
                           window_strides, window_dilation, *,
                           has_only_spatial_dims):
  # Reshape inputs so they are accepted by tf.nn.pool, which expects batch and
  # channel dimensions for operand but not for any of the other inputs.
  if has_only_spatial_dims:  # len(operand_shape) <= 2
    # Call eval_shape on a shape that may contain polynomials, otherwise TF does
    # not know what to do with polynomials in the shape.
    operand_shape = jax2tf._eval_shape(operand_shape)
    # Add batch and channel dimensions to operand.
    operand = tf.reshape(operand, (1,) + operand_shape + (1,))
  else:
    # This branch assumes operand.shape = (batch, spatial_dims, ..., channel),
    # and dimensions, strides, dilation are all (1, spatial_values, ..., 1).
    # Input validation for this is done in _validate_reduce_window_inputs.
    window_dimensions = window_dimensions[1:-1]
    window_strides = window_strides[1:-1]
    window_dilation = window_dilation[1:-1]

  return operand, window_dimensions, window_strides, window_dilation


def _reduce_monoid(operand, window_dimensions, window_strides, padding,
                   base_dilation, window_dilation, computation_name,
                   _in_avals: Sequence[core.ShapedArray],
                   _out_aval: core.ShapedArray):
  dtype = operand.dtype
  # In presence of shape polymorphism, operand.shape may contain None. The
  # actual dimension polynomial shapes are in _in_avals.
  operand_shape = _in_avals[0].shape

  # TODO(marcvanzee): Put reduce_window arguments into dataclass, similar to
  # Gather, to simplify function calls.
  has_only_spatial_dims = _validate_reduce_window_inputs(
      operand_shape, computation_name, dtype, window_dimensions, window_strides,
      base_dilation, window_dilation)

  operand, operand_shape, padding_type = _padding_reduce_window(
      operand, operand_shape, computation_name, window_dimensions,
      window_strides, padding)

  operand, window_dimensions, window_strides, dilations = _reshape_reduce_window(
      operand,
      operand_shape,
      window_dimensions,
      window_strides,
      window_dilation,
      has_only_spatial_dims=has_only_spatial_dims)

  def tf_pool(inputs, pooling_type):
    if any(not core.is_constant_shape(s) for s in
           (window_dimensions, window_strides, dilations)):
      raise NotImplementedError(
          f"TODO: use tf.nn.pool with dynamic shapes¨{window_dimensions=} "
          f" {window_strides=} {dilations=}")
    result = tf.nn.pool(
        inputs,
        window_shape=window_dimensions,
        pooling_type=pooling_type,
        padding=padding_type,
        strides=window_strides,
        dilations=dilations)

    if has_only_spatial_dims:
      # If the input only had spatial dimensions we need to contract the batch
      # and channel dimensions before returning the output.
      result = tf.squeeze(result, [0, -1])

    jax2tf._assert_matching_abstract_shape(result, _out_aval.shape)
    return result

  negate = lambda x: tf.multiply(x, tf.constant(-1, dtype))
  if computation_name == "max":
    return tf_pool(operand, "MAX")
  elif computation_name == "min":
    return negate(tf_pool(negate(operand), "MAX"))
  elif computation_name == "add":
    # TODO(marcvanzee): This may give very large deviations on TPU when using
    # floats as inputs. Alternatively, we could implement this using a
    # convolution with an all-1's kernel.
    return tf.multiply(tf_pool(operand, "AVG"), math.prod(window_dimensions))


def _reduce_window(*args, jaxpr, consts, window_dimensions,
                   window_strides, padding, base_dilation, window_dilation,
                   _in_avals: Sequence[core.ShapedArray],
                   _out_aval: tuple[core.ShapedArray, ...]
                   ) -> tuple[TfVal, ...]:
  assert len(consts) == 0, "Reduction computation cannot have constants"
  operands, init_values = util.split_list(args, [len(args) // 2])

  if len(operands) != 1 or len(init_values) != 1:
    raise _reduce_error("jax2tf does not support variadic reduce_window")

  operand, init_value = operands[0], init_values[0]
  # Infer operation type from jaxpr.
  if (len(jaxpr.eqns) != 1 or
      len(jaxpr.eqns[0].invars) != 2 or
      len(jaxpr.eqns[0].outvars) != 1 or
      jaxpr.eqns[0].primitive.name not in ["min", "max", "add"]):
    raise _reduce_error("Reduction function should be either min, max, or add.")

  computation_name = jaxpr.eqns[0].primitive.name
  result = _reduce_monoid(operand,
                          window_dimensions=window_dimensions,
                          window_strides=window_strides,
                          padding=padding,
                          base_dilation=base_dilation,
                          window_dilation=window_dilation,
                          computation_name=computation_name,
                          _in_avals=(_in_avals[0],),  # Don't pass init_value.
                          _out_aval=_out_aval[0])     # Returns single value.

  reduce_fn = {
      "min": tf.minimum,
      "max": tf.maximum,
      "add": tf.add,
  }[computation_name]
  result = reduce_fn(result, init_value)

  # The outut is expected to be wrapped in a tuple, and since we don't use
  # variadic reductions, this tuple always contains a single element.
  return (result,)


tf_impl_no_xla[lax.reduce_window_min_p] = (
    partial(_reduce_monoid, computation_name="min"))
tf_impl_no_xla[lax.reduce_window_max_p] = (
    partial(_reduce_monoid, computation_name="max"))
tf_impl_no_xla[lax.reduce_window_sum_p] = (
    partial(_reduce_monoid, computation_name="add"))

tf_impl_no_xla[lax.reduce_window_p] = _reduce_window

tf_impl_no_xla[lax.reduce_p] = _unimplemented("reduce")

tf_impl_no_xla[lax.select_and_scatter_add_p] = _unimplemented(
    "select_and_scatter_add")

tf_impl_no_xla[lax.rng_bit_generator_p] = _unimplemented("rng_bit_generator")


def _clip(max_indices: Sequence[TfVal], start_indices: Sequence[TfVal],
          slice_sizes: Sequence[TfVal]):
  """Simulates XLA clipping behavior with TF ops.

  Various TF ops have different clipping behavior than XLA:
  * If `start_indices` is out-of-bounds, then TF fails but XLA clips the indices
  to
    [0, max_len].
  * If `start_indices + slice_size` is out-of-bounds, then TF fails, but XLA
  adjust
    `start_indices` so that a full slice is returned.
  This function clips the start indices correctly.
  """
  # We cast both arguments to `tf.clip_by_value` to int32. Otherwise, this
  # function may return uint32 which is not always compatible with TF ops, so
  # this may result in type errors.
  max_start = tf.cast(tf.subtract(max_indices, slice_sizes), dtype=tf.int32)
  return tf.clip_by_value(tf.cast(start_indices, dtype=tf.int32), 0, max_start)


@dataclasses.dataclass
class GatherArgs:
  operand: TfVal
  start_indices: TfVal
  dnums: lax.GatherDimensionNumbers
  slice_sizes: TfVal
  op_shape: core.Shape
  start_indices_shape: core.Shape
  out_aval: core.ShapedArray

  def __post_init__(self):
    assert len(self.op_shape) == len(self.slice_sizes)

  def __repr__(self):
    return (f"operand shape={self.op_shape}, "
            f"start_indices={self.start_indices}, "
            f"dimension_numbes={self.dnums}, "
            f"slice_sizes={self.slice_sizes}")
  @property
  def batch_dims(self):
    return tuple(x for x in range(len(self.out_aval.shape))
                if x not in self.dnums.offset_dims)

def gather_precondition(precondition_fn: Callable[[GatherArgs], None]):
  """Decorator for specifying a precondition function.

  This decorator should be put on a function with argument `arg` of type
  `GatherArgs`. It will first call `precondition_fn` with `arg` (which may throw
  an exception), and then call the function it is decorating with `arg` as well.
  """

  def decorator(gather_fn: Callable[[GatherArgs], Any]):

    @wraps(gather_fn)
    def wrapper(args: GatherArgs):
      # Call `precondition_fn`; we assume it may throw an exception.
      precondition_fn(args)
      return gather_fn(args)

    return wrapper

  return decorator


def _pre_gather_for_scalar_indexing(args: GatherArgs):
  """Returns True if this call to gather represents scalar indexing into arrays.

  E.g., op[2], op[:, :5, :], jnp.take(op, 0, axis=0).
  """
  # TODO(marcvanzee): Add more assumptions here, because this is currently too
  # permissive.
  if len(args.start_indices_shape) != 1:
    raise ValueError("start_indices shape should be 1")


@gather_precondition(_pre_gather_for_scalar_indexing)
def _gather_for_scalar_indexing(args: GatherArgs):
  """Implements 'scalar indexing into arrays' cases of lax.gather using tf.slice.

  E.g., op[2], op[:, :5, :], jnp.take(op, 0, axis=0).
  """
  indices = tf.expand_dims(args.dnums.start_index_map, 1)
  # lax.gather uses an "index map" which maps `start_indices` to the right axes
  # in `operand`. Since tf.strided_slice uses a single array for specifying the
  # start indices, we use a scatter to map the start indices to the right axes.
  op_shape = jax2tf._eval_shape(args.op_shape)
  slice_sizes_tf = jax2tf._eval_shape(args.slice_sizes)
  # TODO(marcvanzee): Consider transposing `operand`, which is probably more
  # optimization friendly.
  begin = tf.scatter_nd(indices, args.start_indices, [len(op_shape)])
  begin = _clip(op_shape, begin, slice_sizes_tf)
  end = slice_sizes_tf + begin

  # `collapsed_slice_dims` is a tuple of dimensions to collapse, e.g. (0, 2).
  # `tf.strided_slice` expects a binary mask to specify the shrink axes, i.e.,
  # if we want to shrink axis 0 and 2, this corresponds to binary mask 101,
  # which is 5 in decimals. The following line converts the lax representation
  # to the one used by `tf.strided_slice`.
  shrink_mask = sum(2**x for x in args.dnums.collapsed_slice_dims)
  res = tf.strided_slice(args.operand, begin, end, shrink_axis_mask=shrink_mask)
  # Shape inference doesn't work for tf.strided_slice.
  res = jax2tf._ensure_tf_shape_if_dynamic(
      res, jax2tf._aval_to_tf_shape(args.out_aval)
  )
  return res


def _pre_gather_for_multidim_indexing(args: GatherArgs):
  """Returns True if this call to gather represents multi-dimensional indexing.

  E.g., jnp.take(op, [[0], [1]], axis=0).
  Note we currently only support multi-dimensional indexing if the last
  dimension is 1.
  """
  # Handle only the case when tf.gather argument batch_dims=0.
  # Find axis to match the tf.gather semantics
  # Let I = len(start_indices_shape)
  # let O = len(op_shape)
  # slice_sizes == op_shape[:axis] + (1,) + op_shape[axis+1:]
  # collapsed_slice_dims == (axis,)
  # start_index_map == (axis,)
  # offset_dims == (0, 1, ..., axis - 1, axis + I, ..., O + I - 1)
  # We added a trailing dimension of size 1
  op_shape = args.op_shape
  start_index_map = args.dnums.start_index_map
  collapsed_slice_dims = args.dnums.collapsed_slice_dims
  offset_dims = args.dnums.offset_dims
  if not (len(op_shape) >= 1 and len(start_index_map) == 1 and
          len(collapsed_slice_dims) == 1 and collapsed_slice_dims[0]
          == start_index_map[0] and len(offset_dims) == len(op_shape) - 1):
    raise ValueError("unsupported dimension numbers")
  # We added a trailing dimension of size 1
  if not core.definitely_equal(args.start_indices_shape[-1], 1):
    raise ValueError("start_indices shape[-1] should be 1")
  # Guess the axis
  axis = collapsed_slice_dims[0]
  index_dims = len(args.start_indices_shape) - 1
  expected_offset_dims = tuple(
      list(range(axis)) +
      list(range(axis + index_dims,
                 len(op_shape) + index_dims - 1)))
  if offset_dims != expected_offset_dims:
    raise ValueError("unsupported offset_dims")
  expected_slice_sizes = op_shape[:axis] + (1,) + op_shape[axis + 1:]  # type: ignore
  if not core.definitely_equal_shape(args.slice_sizes, expected_slice_sizes):
    raise ValueError("unsupported slice_sizes")


@gather_precondition(_pre_gather_for_multidim_indexing)
def _gather_for_multidim_indexing(args: GatherArgs):
  """Implements 'multi-dimensional indexing into arrays' cases of lax.gather using tf.gather.

  E.g., jnp.take(op, [[0], [1]], axis=0).
  """
  # Guess the axis.
  axis = args.dnums.collapsed_slice_dims[0]
  squeezed_indices = tf.squeeze(args.start_indices, -1)
  op_shape = jax2tf._eval_shape(args.op_shape)
  start_indices = _clip((op_shape[axis],), squeezed_indices, (1,))
  return tf.gather(args.operand, start_indices, axis=axis, batch_dims=0)


def _pre_gather_with_batch_dim(args: GatherArgs):
  """Returns True if this call to gather has non-empty batch dimensions.

  This is for instance triggered when doing jax.vmap(lax.dynamic_slice).
  """
  # We assume exactly one batch (and one or more non-batch dimensions).
  if len(args.batch_dims) != 1:
    raise ValueError(f"batch_dims is {len(args.batch_dims)} but should be 1")

  # `start_index_map` maps indices in `start_indices` to indices in `operand`.
  # For simplicity, we currently only consider the case where this mapping is
  # the identity function, i.e., [2, 3] in `start_indices` maps to
  # `operand[2, 3]`.
  if args.dnums.start_index_map != tuple(range(args.start_indices_shape[-1])):
    raise ValueError("unsupported start_index_map")

  # The batch dims in `start_indices` and `operand` should agree.
  if not core.definitely_equal(args.op_shape[0], args.start_indices_shape[0]):
    raise ValueError("Batch dimensions in operand and start_indices don't "
                     "agree")


def _pre_gather_with_batch_dims(args: GatherArgs):
  """Returns True if this call to gather has non-empty 2D batch dimensions.

  This is for instance triggered when doing
  jax.vmap(jax.vmap(lax.dynamic_slice)).
  """
  if len(args.dnums.collapsed_slice_dims) != 0:
    # NOTE: this can be relaxed in _gather_with_batch_dims but we might
    #   also need to re-work the output reshaping
    raise ValueError("only len(collapsed_slice_dims) == 0 is supported")

  # NOTE: This supports higher dimensions than listed (the highest dimenison
  # in the tests is 3D so it is limited to that, but the implementation is
  # designed to handle higher dimensions (N-Dimensional)).
  if len(args.batch_dims) not in [1, 2, 3]:
    raise ValueError(
        f"Size of batch_dims is {len(args.batch_dims)} but should be up to 3"
    )

@gather_precondition(_pre_gather_with_batch_dim)
def _gather_with_batch_dim(args: GatherArgs):
  """Implements call to gather with non-empty batch dimensions.

  E.g., when doing `jax.vmap(lax.dynamic_slice).
  """
  op_shape = jax2tf._eval_shape(args.op_shape)
  start_indices = _clip(op_shape, args.start_indices, args.slice_sizes)
  result = tf.map_fn(
    lambda idxs: tf.slice(args.operand, begin=idxs, size=args.slice_sizes),
    start_indices,
    fn_output_signature=jax2tf._to_tf_dtype(args.operand.dtype)
  )
  result = tf.reshape(result, jax2tf._eval_shape(args.out_aval.shape))
  return result


def _gather_generate_indices(shape: tuple[int, ...]):
  """
  Returns the indices of the according to `shape`:
    each element in the output is the index of an element of an array
    of the provided shape. The result's shape is (math.prod(shape), len(shape))

  For example, given shape (2,2) it returns (0,0),(0,1),(1,0),(1,1)
  """
  return tf.reshape(
      tf.stack(
          tf.meshgrid(
              *[tf.range(start=0, limit=x) for x in shape], indexing="ij"
          ),
          axis=-1,
      ),
      (-1, len(shape)),
  )


@gather_precondition(_pre_gather_with_batch_dims)
def _gather_with_batch_dims(args: GatherArgs):
  """Implements call to gather with non-empty 2D batch dimensions."""
  op_shape = jax2tf._eval_shape(args.op_shape)
  output_shape = jax2tf._eval_shape(args.out_aval.shape)
  # Used to map the start_indices w.r.t start_index_map
  indices = tf.expand_dims(args.dnums.start_index_map, 1)

  # batch_indices is shaped (N,d) where N is the number of slices and d is
  # the number of batch_dims; batch_indices_size equals to N
  batch_indices = _gather_generate_indices(
      tuple(output_shape[i] for i in args.batch_dims)
  )
  batch_indices_size = jax2tf._eval_shape(batch_indices.shape)[0]
  # offset_indices is shaped (K,d) where K is the number of elements in each
  # slice and d is the number of offset_dims; offset_indices_size equals to K
  offset_indices = _gather_generate_indices(
      tuple(output_shape[i] for i in args.dnums.offset_dims)
  )
  offset_indices_size = jax2tf._eval_shape(offset_indices.shape)[0]

  # After we compute the result we need to reshape the axes with respect to
  # the output batch_dims and offset_dims.
  dim_mask = args.batch_dims + args.dnums.offset_dims
  mask_output_shape = tuple(output_shape[x] for x in dim_mask)

  def get_scatter_indices(indices, batch_indices_size, size_of_index_map):
    """Generate the start indices of each slice, which index into the operand."""
    # Tile indices batch_indices_size times
    tiled_indices = tf.tile(
        tf.expand_dims(indices, 0), [batch_indices_size, 1, 1]
    )
    # The above tiles need to index the proper element of batch_indices
    # To do this generate a repeated sequence of numbers
    temp_batch_indices = tf.repeat(
        tf.range(start=0, limit=batch_indices_size), size_of_index_map
    )
    # Reshape the above sequence so it follows the same shape of tiled_indices
    temp_batch_indices = tf.reshape(
        temp_batch_indices, (batch_indices_size, size_of_index_map, 1)
    )
    # Now we concatenate to create indices offset by the temp_batch_indices
    return tf.concat([temp_batch_indices, tiled_indices], axis=-1)

  slice_start_indices = tf.gather_nd(args.start_indices, batch_indices)
  # TODO: In the case where start_index_map is the identity we can skip this.
  scatter_indices = get_scatter_indices(
      indices, batch_indices_size, len(args.dnums.start_index_map)
  )
  # We map the scatter_indices w.r.t start_index_map
  indices_in_operand = tf.scatter_nd(
        scatter_indices, slice_start_indices, [batch_indices_size, len(op_shape)]
  )

  # We clip the indices as OOB cases are possible when offsetting past
  # the operand boundaries
  clipped_start_indices = _clip(op_shape, indices_in_operand, args.slice_sizes)
  # Here we need to broadcast clipped_start_indices and add each of the offsets
  # which will generate a large index tensor of shape (T,d) where T is the
  # number of slices times the size of each slice (i.e total number of items
  # across all sices); d is rank(operand)
  slice_element_indices = tf.add(
      tf.repeat(clipped_start_indices, offset_indices_size, axis=0),
      tf.tile(offset_indices, (batch_indices_size, 1)),
  )
  results = tf.gather_nd(args.operand, slice_element_indices)

  # Here results comes shaped as (N,1). Because collapsed_slice_dims is 0,
  # offset_dims is effectviely slice_sizes.
  # We reshape to mask_output_shape because if we directly reshape to the
  # output shape and our batch_dims are non-contiguous we will produce the
  # wrong shape. Reshaping to mask_output_shape gives (...,*slice_sizes),
  # which we then transpose to permute the axes in the proper way.
  # Note that if the batch_dims are contiguous this won't change the output.
  temp = tf.reshape(results, shape=mask_output_shape)
  return tf.transpose(temp, perm=tf.math.invert_permutation(dim_mask))

def _gather(operand, start_indices, *, dimension_numbers,
            slice_sizes: core.Shape, indices_are_sorted, unique_indices, mode,
            fill_value, _in_avals: Sequence[core.ShapedArray],
            _out_aval: core.ShapedArray):
  """Tensorflow implementation of gather."""
  if mode == lax.GatherScatterMode.FILL_OR_DROP:
    gather_fill_fn = jax2tf._convert_jax_impl(lax_slicing._gather_fill,
                                              multiple_results=False)
    return gather_fill_fn(
        operand, start_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=_out_aval.shape, _in_avals=_in_avals, _out_aval=_out_aval)

  # TODO(marcvanzee): Check if we need more tests in shape_poly for gather with
  # enable_xla=False.
  gather_args = GatherArgs(
      operand=operand,
      start_indices=start_indices,
      dnums=dimension_numbers,
      slice_sizes=slice_sizes,
      op_shape=_in_avals[0].shape,
      start_indices_shape=_in_avals[1].shape,
      out_aval=_out_aval)

  errors = []

  for gather_fn in [
      _gather_for_scalar_indexing,
      _gather_for_multidim_indexing,
      _gather_with_batch_dim,
      _gather_with_batch_dims,
  ]:
    try:
      return gather_fn(gather_args)
    except ValueError as e:
      errors.append(f"{gather_fn}: {repr(e)}")

  error_msg = (f"Unsupported arguments for gather: {gather_args}, errors:\n" +
               "\n".join(errors))

  raise _error("gather", error_msg)


tf_impl_no_xla[lax.gather_p] = _gather


def _dynamic_slice(operand, *start_indices, slice_sizes: core.Shape,
                   _in_avals: Sequence[core.ShapedArray],
                   _out_aval: core.ShapedArray):
  start_indices = tf.stack(start_indices)
  slice_sizes_tf = jax2tf._eval_shape(slice_sizes)

  operand_shape = jax2tf._eval_shape(_in_avals[0].shape)
  start_indices = _clip(operand_shape, start_indices, slice_sizes_tf)
  return tf.slice(operand, start_indices, size=slice_sizes_tf)


tf_impl_no_xla[lax.dynamic_slice_p] = _dynamic_slice


def _dynamic_update_slice(operand, update, *start_indices,
                          _in_avals: Sequence[core.ShapedArray],
                          _out_aval: core.ShapedArray):
  start_indices = tf.stack(start_indices)

  op_shape = jax2tf._eval_shape(_in_avals[0].shape)
  op_size = tf.size(operand)
  update_shape_tf = jax2tf._eval_shape(_in_avals[1].shape)

  start_indices = _clip(op_shape, start_indices, update_shape_tf)
  end_indices = tf.add(start_indices, update_shape_tf)

  # Get the cells to update in `operand` as an array of ids.
  id_tensor = tf.reshape(tf.range(op_size), op_shape)
  scattered_indices = tf.strided_slice(id_tensor, start_indices, end_indices)

  # Create an array containing updates at scattered_indices and zeros otherwise.
  flat_indices = tf.expand_dims(tf.nest.flatten(scattered_indices), -1)
  flat_update = tf.nest.flatten(update)
  update = tf.scatter_nd(flat_indices, flat_update, (op_size,))
  update = tf.reshape(update, op_shape)

  # Create a bool mask that is True only where `operand` should be updated.
  update_mask = tf.ones_like(flat_update, dtype=tf.bool)
  update_mask = tf.scatter_nd(flat_indices, update_mask, (op_size,))
  update_mask = tf.reshape(update_mask, op_shape)

  # Use the mask to only update `operand` with `update`.
  return tf.where(update_mask, update, operand)


tf_impl_no_xla[lax.dynamic_update_slice_p] = _dynamic_update_slice


def shift_axes_forward(operand,
                       axes: tuple[int, ...],
                       inverse: bool = False,
                       forward: bool = True):
  """Shifts the tuple of axes to the front of an array"""
  other_axes = tuple([i for i in range(len(operand.shape)) if i not in axes])
  fwd_order = axes + other_axes if forward else other_axes + axes
  order = fwd_order if not inverse else _invert_permutation(fwd_order)
  return tf.transpose(operand, order)

def convert_scatter_jax_to_tf(update_op, unsorted_segment_op=None):

  def _sparse_scatter(operand, scatter_indices, updates, unique_indices, mode,
                      _in_avals: Sequence[core.ShapedArray],
                      _out_aval: core.ShapedArray):
    """Implementation of scatter specialised to indexing from the front axes.

    This covers unique indices and non-unique indices of single depth.
    Note on unique indices: `tf.tensor_scatter_nd_update` interprets indices
    thusly: every axis except the final one encodes a batch dimension, the final
    axis encoding the actual indices to scatter in to. It enforces, at least
    one, batch dimension so we add an empty dimension to indices and updates if
    lacking.

    Note on non-unique indices: There is no tf op for non-single depth indexing,
    but if indexing is single depth, this can be viewed as a segment op.
    """
    # Infer unique indices from lack of batch dimension
    unique_indices = unique_indices or (len(scatter_indices.shape) == 1)
    if unique_indices:
      suboperand = tf.gather_nd(operand, scatter_indices)
      updated_suboperand = update_op(suboperand, updates)
      # add a batch dim if none exist
      if len(scatter_indices.shape) == 1:
        scatter_indices = scatter_indices[None]
        updated_suboperand = updated_suboperand[None]
      y = tf.tensor_scatter_nd_update(operand, scatter_indices, updated_suboperand)
    else:
      if (scatter_indices.shape[-1] == 1) and unsorted_segment_op:
        # If only indexing into the first dimension, it's a segment op
        operand_update = unsorted_segment_op(updates,
                                             tf.squeeze(scatter_indices, -1),
                                             operand.shape[0])
        y = update_op(operand, operand_update)
      else:
        raise _scatter_error(
            "Scatter only supports non-unique "
            "indices with indexing into only one dimension for (add, mul, min, "
            "max)")
    return y

  def sparse_scatter(operand, scatter_indices, updates, update_jaxpr,
                     update_consts, dimension_numbers, indices_are_sorted: bool,
                     unique_indices: bool, mode,
                     _in_avals: Sequence[core.ShapedArray],
                     _out_aval: core.ShapedArray):
    """
    Wrapper around the scatter function.
    The underlying tf ops `tf.tensor_scatter_nd_update` and
    `tf.math.unsorted_segment_*` index from the front dimensions.
    `tf.math.unsorted_segment_*` indexs to a depth 1 from the front.
    `tf.tensor_scatter_nd_update` indexs from the front dimensions onwards,
    with no ability to skip a dimension. This function shifts the axes to be
    indexed to the front then calls a front-specific implementation, then
    inverse-shifts the output.

    scatter_dims_to_operand_dims: dimensions which the scatter indexes in to.
      We shift these to the front to match tf syntax. All other dims are batch
    update_window_dims: dimensions which are not batch dimensions. We shift
      these to the back as the remaining dimensions are batch dimensions.
    """
    del update_jaxpr, update_consts, indices_are_sorted  # Unused arguments

    update_window_dims = dimension_numbers.update_window_dims
    inserted_window_dims = dimension_numbers.inserted_window_dims
    scatter_to_operand_dims = dimension_numbers.scatter_dims_to_operand_dims

    dtype = operand.dtype  # assume updates has same dtype as operand
    if dtype in [tf.bool, tf.complex64]:
      raise _scatter_error(f"Scatter does not support operands of type {dtype}")

    if inserted_window_dims != scatter_to_operand_dims:
      raise _scatter_error("Complex scatters are not supported")

    if (mode != lax.GatherScatterMode.FILL_OR_DROP and
        mode != lax.GatherScatterMode.PROMISE_IN_BOUNDS):
      # The OOB behavior for tf.scatter is as follows:
      # - When running in eager or graph mode, it throws an error.
      #   TODO(marcvanzee): Fix this case by removing the OOB indices.
      # - When running in compile mode, the OOB indices are dropped, which is
      #   the same behavior as FILL_OR_DROP and PROMISE_IN_BOUNDS.
      # To ensure correctness, we disallow CLIP mode for now.
      raise _scatter_error("Only scatter modes `FILL_OR_DROP` and "
                           "`PROMISE_IN_BOUNDS` are supported.")

    # Shift axes to the front to match tf syntax, inverse afterwards
    fwd = partial(shift_axes_forward, axes=scatter_to_operand_dims)
    inv = partial(fwd, inverse=True)

    # Shift update value axes to the back, so batch are at the front
    updates_shifted = shift_axes_forward(
        updates, axes=update_window_dims, forward=False)

    return inv(
        _sparse_scatter(
            fwd(operand), scatter_indices, updates_shifted, unique_indices,
            mode, _in_avals, _out_aval))
  return sparse_scatter


tf_impl_no_xla[lax.scatter_p] = convert_scatter_jax_to_tf(
    lambda x, y: y)  # just replace with the update
tf_impl_no_xla[lax.scatter_add_p] = convert_scatter_jax_to_tf(tf.add,      tf.math.unsorted_segment_sum)
tf_impl_no_xla[lax.scatter_mul_p] = convert_scatter_jax_to_tf(tf.multiply, tf.math.unsorted_segment_prod)
tf_impl_no_xla[lax.scatter_min_p] = convert_scatter_jax_to_tf(tf.minimum,  tf.math.unsorted_segment_min)
tf_impl_no_xla[lax.scatter_max_p] = convert_scatter_jax_to_tf(tf.maximum,  tf.math.unsorted_segment_max)

tf_impl_no_xla[lax.sort_p] = _unimplemented("sort")

tf_impl_no_xla[lax.reduce_precision_p] = _unimplemented("reduce_precision")
