Repository: scikit-learn/scikit-learn
File: sklearn/metrics/_classification.py (Python)

Code before:

    # Finally, we have all our sufficient statistics. Divide! #
    max_float_type = _max_precision_float_dtype(xp=xp, device=device_)
    beta2 = xp.asarray(beta**2, dtype=max_float_type, device=device_)

    # Divide, and on zero-division, set scores and/or warn according to
    # zero_division:
    precision = _prf_divide(
        tp_sum, pred_sum, "precision", "predicted", average, warn_for, zero_division
    )
    recall = _prf_divide(
        tp_sum, true_sum, "recall", "true", average, warn_for, zero_division
    )

    if np.isposinf(beta):
        f_score = recall
    elif beta == 0:
        f_score = precision
    else:
        # The score is defined as:
        # score = (1 + beta**2) * precision * recall / (beta**2 * precision + recall)
        # Therefore, we can express the score in terms of confusion matrix entries as:
        # score = (1 + beta**2) * tp / ((1 + beta**2) * tp + beta**2 * fn + fp)
        denom = beta2 * xp.asarray(
            true_sum, dtype=max_float_type, device=device_
        ) + xp.asarray(pred_sum, dtype=max_float_type, device=device_)
        f_score = _prf_divide(
            (1 + beta2) * xp.asarray(tp_sum, dtype=max_float_type, device=device_),
            denom,
            "f-score",
            "true nor predicted",
            average,
            warn_for,
            zero_division,
        )

    # Average the results
    if average == "weighted":
        weights = true_sum
    elif average == "samples":
        weights = sample_weight
    else:
        weights = None

    if average is not None:
        precision = float(_nanaverage(precision, weights=weights))
        recall = float(_nanaverage(recall, weights=weights))
        f_score = float(_nanaverage(f_score, weights=weights))
        true_sum = None  # return no support

    return precision, recall, f_score, true_sum

Review comment (suggestion):
shouldn't `true_sum` and `pred_sum` already be on the device? I don't like that a simple multiplication and sum has become such a long piece of code.

Code after:

    # Finally, we have all our sufficient statistics. Divide! #
    beta2 = beta**2

    # Divide, and on zero-division, set scores and/or warn according to
    # zero_division:
    precision = _prf_divide(
        tp_sum, pred_sum, "precision", "predicted", average, warn_for, zero_division
    )
    recall = _prf_divide(
        tp_sum, true_sum, "recall", "true", average, warn_for, zero_division
    )

    if np.isposinf(beta):
        f_score = recall
    elif beta == 0:
        f_score = precision
    else:
        # The score is defined as:
        # score = (1 + beta**2) * precision * recall / (beta**2 * precision + recall)
        # Therefore, we can express the score in terms of confusion matrix entries as:
        # score = (1 + beta**2) * tp / ((1 + beta**2) * tp + beta**2 * fn + fp)

        # Array api strict requires all arrays to be of the same type so we
        # need to convert true_sum, pred_sum and tp_sum to the max supported
        # float dtype because beta2 is a float
        max_float_type = _max_precision_float_dtype(xp=xp, device=device_)
        denom = beta2 * xp.astype(true_sum, max_float_type) + xp.astype(
            pred_sum, max_float_type
        )
        f_score = _prf_divide(
            (1 + beta2) * xp.astype(tp_sum, max_float_type),
            denom,
            "f-score",
            "true nor predicted",
            average,
            warn_for,
            zero_division,
        )

    # Average the results
    if average == "weighted":
        weights = true_sum
    elif average == "samples":
        weights = sample_weight
    else:
        weights = None

    if average is not None:
        precision = float(_nanaverage(precision, weights=weights))
        recall = float(_nanaverage(recall, weights=weights))