code: 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)
      )
  
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 34 for calling a global function with example usage: _gather_with_batch_dims(args) and returns: tf, tf