def predict_logits(img_x, phrases_x, f_similarity):
    """
    Compute given similarity measure over the last dimension (features) of `img_x` and `phrases_x`.

    :param img_x: A [*, d1, d2, d3] tensor
    :param phrases_x: A [*, d1, d2, d3] tensor
    :param f_similarity: A similarity measure between [*, d3] tensors
    :return: A `[*, d1, d2]` tensor
    """
    similarity = f_similarity(img_x, phrases_x, dim=-1)

    return similarity