import torch


def flat_to_batch(source: torch.Tensor, scope: torch.Tensor) -> torch.Tensor:
    """
    :param source: A tensor of shape (num_bonds, hidden_size) containing message features.
    :param scope: A tensor of shape (batch_size, max_num_bonds) expressing bond indices for each mol/row.
    :return: A tensor of shape (batch, max_num_bonds, hidden_size) containing the message features.
    """
    final_size = (scope.shape[0], -1,
                  source.shape[1])  # batch x max_num_bonds x hidden_size
    ret = source.index_select(dim=0, index=scope.view(-1))
    return ret.view(final_size)