import torch


def scatter_embs(input_embs, inputs):
    """
    For inputs that have 'input_embs' field passed in, replace the entry in input_embs[i] with the entry
    from inputs[i]['input_embs']. This is useful for the Integrated Gradients - for which the predict is
    called with inputs with 'input_embs' field which is an interpolation between the baseline and the real calculated
    input embeddings for the sample.
    :param input_embs: tensor of shape B x S x h of input embeddings according to the input sentences.
    :param inputs: list of dictionaries (smaples), for which the 'input_embs' field might be specified
    :return: tensor of shape B x S x h with embeddings (if passed) from inputs inserted to input_embs
    """
    interp_embeds = [(ind, ex.get('input_embs')) for ind, ex in enumerate(inputs)]
    for ind, embed in interp_embeds:
        if embed is not None:
            input_embs[ind] = torch.tensor(embed)

    return input_embs