See https://gitlab.cern.ch/atlas-flavor-tagging-tools/algorithms/GNNJetTagger/-/blob/main/gnn_tagger/training/metrics.py