Add AUROC metric option to ClassificationTask.

[favyen2]

No description provided.

[favyen2]

@claude can you review this PR?

[claude]

Claude finished @favyen2's task in 1m 13s —— View job

---

PR Review: Add AUROC metric option to ClassificationTask

• Gather context and read changed files
• Review code for correctness, style, and best practices
• Post review feedback

---

Overall

Clean, minimal addition that follows the established pattern for optional metrics in ClassificationTask. The implementation is functionally correct.

Positive observations

• Import is in alphabetical order (consistent with surrounding imports).
• Constructor parameters, instance attribute storage, and docstring entries all follow the exact same pattern as enable_f1_metric / f1_metric_kwargs.
• num_classes is correctly injected before user-supplied auroc_metric_kwargs, so users can't accidentally override it to something wrong (same pattern as accuracy and F1).
• MulticlassAUROC receives probability scores from ClassificationHead (after softmax), so the metric inputs are semantically correct.

Issues / suggestions

1. No unit test (minor)

There's a clear existing pattern in tests/unit/train/tasks/test_classification.py — test_per_class_f1() verifies the F1 metric end-to-end using get_metrics(). A similar test for AUROC (e.g. two classes, a few probability vectors with known ground truth → assert the AUROC value) would give confidence the wiring is right and prevent regression. Fix this →

2. No per-class AUROC (informational)

The F1 metric has special handling when average=None — it generates per-class metric keys ({class_name}_f1, etc.) and wraps each in a ClassificationMetric(metric, class_id). The AUROC implementation has no equivalent. If a user passes average=None via auroc_metric_kwargs, ClassificationMetric.compute() will return a raw tensor (one value per class) instead of a scalar, which may behave unexpectedly in logging. This is worth either:

• Documenting in the docstring that per-class AUROC is not supported (must use average="macro" or "weighted"), or
• Adding the same per-class expansion as F1.

Neither is blocking for this PR, but worth noting.

Verdict

LGTM with the above notes. The implementation is correct and consistent. Adding a unit test would be ideal before merging.

---