def unflatten_beam_dim(tensor, batch_size, num_beams):
            """Unflattens the first, flat batch*beam dimension of a non-scalar array."""
            # ignore scalars (e.g. cache index)
            if tensor.ndim == 0:
                return tensor
            return tensor.reshape((batch_size, num_beams) + tensor.shape[1:])
Unflattens the first, flat batch*beam dimension of a non-scalar array.
dpkd/transformers/src/transformers/generation/flax_utils.py