code: def _tensor_gather_helper(
      gather_indices,
      gather_from,
      batch_size,
      range_size,
      gather_shape,
      name=None,
  ):
      """Helper for gathering the right indices from the tensor.
  
      This works by reshaping gather_from to gather_shape (e.g. [-1]) and then
      gathering from that according to the gather_indices, which are offset by
      the right amounts in order to preserve the batch order.
  
      Args:
        gather_indices: The tensor indices that we use to gather.
        gather_from: The tensor that we are gathering from.
        batch_size: The input batch size.
        range_size: The number of values in each range. Likely equal to
          beam_width.
        gather_shape: What we should reshape gather_from to in order to preserve
          the correct values. An example is when gather_from is the attention
          from an AttentionWrapperState with shape
          [batch_size, beam_width, attention_size]. There, we want to preserve
          the attention_size elements, so gather_shape is
          [batch_size * beam_width, -1]. Then, upon reshape, we still have the
          attention_size as desired.
        name: The tensor name for set of operations. By default this is
          'tensor_gather_helper'. The final output is named 'output'.
  
      Returns:
        output: Gathered tensor of shape
          tf.shape(gather_from)[:1+len(gather_shape)]
      """
      with tf.name_scope(name or "tensor_gather_helper"):
          range_ = tf.expand_dims(tf.range(batch_size) * range_size, 1)
          gather_indices = tf.reshape(gather_indices + range_, [-1])
          output = tf.gather(
              tf.reshape(gather_from, gather_shape), gather_indices
          )
          final_shape = tf.shape(gather_from)[: 1 + len(gather_shape)]
          static_batch_size = tf.get_static_value(batch_size)
          final_static_shape = tf.TensorShape(
              [static_batch_size]
          ).concatenate(gather_from.shape[1 : 1 + len(gather_shape)])
          output = tf.reshape(output, final_shape, name="output")
          output.set_shape(final_static_shape)
          return output
  
imports: import collections
  import numpy
  import tensorflow
  from tensorflow_addons import options
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import decoder
  from tensorflow_addons.utils import keras_utils
  from tensorflow_addons.utils.resource_loader import LazySO
  from tensorflow_addons.utils.types import FloatTensorLike, TensorLike, Number
  from typeguard import typechecked
  from typing import Callable, Optional
label_desc: Use image node_id 17 for calling a global function with example usage: _tensor_gather_helper(gather_indices, gather_from, batch_size, range_size, gather_shape, name) and returns: output