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