code: def _check_batch_beam(t, batch_size, beam_width):
      """Returns an Assert operation checking that the elements of the stacked
      TensorArray can be reshaped to [batch_size, beam_size, -1].
  
      At this point, the TensorArray elements have a known rank of at
      least 1.
      """
      error_message = (
          "TensorArray reordering expects elements to be "
          "reshapable to [batch_size, beam_size, -1] which is "
          "incompatible with the dynamic shape of %s elements. "
          "Consider setting reorder_tensor_arrays to False to disable "
          "TensorArray reordering during the beam search."
          % (t if tf.executing_eagerly() else t.name)
      )
      rank = t.shape.ndims
      shape = tf.shape(t)
      if rank == 2:
          condition = tf.equal(shape[1], batch_size * beam_width)
      else:
          condition = tf.logical_or(
              tf.equal(shape[1], batch_size * beam_width),
              tf.logical_and(
                  tf.equal(shape[1], batch_size),
                  tf.equal(shape[2], beam_width),
              ),
          )
      return tf.Assert(condition, [error_message])
  
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 8 for calling a global function with example usage: _check_batch_beam(t, batch_size, beam_width) and returns: tf