code: 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")
  
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 28 for calling a global function with example usage: _pre_gather_for_multidim_indexing(args) without return types