def group_att_over_input(inp_att_vals, inp_att_mask, att_indxs, att_indxs_mask,
                         inp_att_keys=None):
    """
    Returns properly concatenated and shaped attention keys/values/mask over
    input (e.g. hidden states and embds).

    :param inp_att_keys: [inp_batch_size, seq_len, att_hidden_dim]
    :param inp_att_vals: [inp_batch_size, seq_len, val_dim]
    :param inp_att_mask: [inp_batch_size, seq_len]
    :param att_indxs: [out_batch_size, els_per_unit]
    :param att_indxs_mask: [out_batch_size, els_per_unit]
    :return: att_keys: [out_batch_size, total_els, att_hidden_dim]
             att_values: [out_batch_size, total_els, val_dim]
             att_mask: [out_batch_size, total_els]
    """
    seq_len = inp_att_vals.size(1)
    out_batch_size = att_indxs.size(0)
    els_per_unit = att_indxs.size(1)
    total_els_per_unit = els_per_unit * seq_len

    inp_att_vals = inp_att_vals[att_indxs] * \
                   att_indxs_mask.unsqueeze(-1).unsqueeze(-1)
    inp_att_mask = inp_att_mask[att_indxs] * \
                   att_indxs_mask.unsqueeze(-1)

    att_vals = inp_att_vals.view(out_batch_size, total_els_per_unit, -1)
    att_mask = inp_att_mask.view(out_batch_size, total_els_per_unit)

    if inp_att_keys is not None:
        inp_att_keys = inp_att_keys[att_indxs] * \
                       att_indxs_mask.unsqueeze(-1).unsqueeze(-1)
        att_keys = inp_att_keys.view(out_batch_size, total_els_per_unit, -1)
        return att_keys, att_vals, att_mask
    else:
        return att_vals, att_mask