def extend_enc_output(tensor, num_beams=None):
            """
            Broadcast tensor with `num_beams` replica, with correct order Input: tensor of shape (batch_size*n_docs ,
            d) Output: tensor of shape (batch_size*num_beams*n_docs , d)
            """

            # expand batch_size & num_beam dimensions
            d_shape_list = tensor.shape[1:]

            # split n_docs dimensions
            new_shape = (batch_size, 1, n_docs) + d_shape_list
            tensor = tf.reshape(tensor, new_shape)

            # repeat same last hidden states over `num_beams` dimension
            new_shape = (batch_size, num_beams, n_docs) + d_shape_list
            tensor = tf.broadcast_to(tensor, new_shape)

            # merge `batch_size`, `num_beams`, `num_docs` dims again
            new_shape = (batch_size * num_beams * n_docs,) + d_shape_list
            return tf.reshape(tensor, new_shape)

            Broadcast tensor with `num_beams` replica, with correct order Input: tensor of shape (batch_size*n_docs ,
            d) Output: tensor of shape (batch_size*num_beams*n_docs , d)
            
dpkd/transformers/src/transformers/models/rag/modeling_tf_rag.py