def _do_test_pimo_outputs(
    thresholds: Tensor,
    shared_fpr: Tensor,
    per_image_tprs: Tensor,
    image_classes: Tensor,
    expected_thresholds: Tensor,
    expected_shared_fpr: Tensor,
    expected_per_image_tprs: Tensor,
    expected_image_classes: Tensor,
) -> None:
    """Test if the outputs of any of the PIMO interfaces are correct."""
    assert isinstance(shared_fpr, Tensor)
    assert isinstance(per_image_tprs, Tensor)
    assert isinstance(image_classes, Tensor)
    assert isinstance(expected_thresholds, Tensor)
    assert isinstance(expected_shared_fpr, Tensor)
    assert isinstance(expected_per_image_tprs, Tensor)
    assert isinstance(expected_image_classes, Tensor)
    allclose = torch.allclose

    assert thresholds.ndim == 1
    assert shared_fpr.ndim == 1
    assert per_image_tprs.ndim == 2
    assert tuple(image_classes.shape) == (3,)

    assert allclose(thresholds, expected_thresholds)
    assert allclose(shared_fpr, expected_shared_fpr)
    assert allclose(per_image_tprs, expected_per_image_tprs, equal_nan=True)
    assert (image_classes == expected_image_classes).all()
Test if the outputs of any of the PIMO interfaces are correct.
tests/unit/metrics/pimo/test_pimo.py