src/transformers/generation/flax_utils.py/FlaxGenerationMixin/_expand_to_num_beams
class FlaxGenerationMixin:    def _expand_to_num_beams(tensor, num_beams):
        return jnp.broadcast_to(tensor[:, None], (tensor.shape[0], num_beams) + tensor.shape[1:])