def METHOD_NAME(batch_attention_mask):
    """Pad attention_mask."""
    batch_size = len(batch_attention_mask)
    max_len = max(map(len, batch_attention_mask))
    attention_mask = np.ones((batch_size, max_len, max_len), dtype="float32") * -1e9
    for i, mask_data in enumerate(attention_mask):
        seq_len = len(batch_attention_mask[i])
        if pad_right:
            mask_data[:seq_len:, :seq_len] = np.array(batch_attention_mask[i], dtype="float32")
        else:
            mask_data[-seq_len:, -seq_len:] = np.array(batch_attention_mask[i], dtype="float32")
    # In order to ensure the correct broadcasting mechanism, expand one
    # dimension to the second dimension (n_head of Transformer).
    attention_mask = np.expand_dims(attention_mask, axis=1)
    return attention_mask