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