code: 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_*` indexes to a depth 1 from the front.
          `tf.tensor_scatter_nd_update` indexes 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
  
imports: from __future__ import annotations
  import builtins
  from collections.abc import Sequence
  import dataclasses
  from functools import partial, wraps
  import math
  import string
  from typing import Any, Callable, Optional
  from jax._src import core
  from jax import lax
  from jax._src.lax import slicing
  from jax._src import dtypes
  from jax._src import util
  from jax.experimental.jax2tf import jax2tf
  import numpy
  import tensorflow
label_desc: Use image node_id 39 for calling a global function with example usage: convert_scatter_jax_to_tf(update_op, unsorted_segment_op) and returns: sparse_scatter, y, inv