def unflatten_beam_dim(tensor):
                """Unflattens the first, flat batch*beam dimension of a non-scalar array."""
                shape = shape_list(tensor)
                return tf.reshape(tensor, [-1, generation_config.num_beams] + shape[1:])
Unflattens the first, flat batch*beam dimension of a non-scalar array.
dpkd/transformers/src/transformers/models/rag/modeling_tf_rag.py