[sylvesterkaczmarek]

Problem

Marginal, joint, and polyadic classification likelihoods materialize softmax probabilities before taking their logarithm. Finite logits such as [0, -1000] therefore produce infinite negative log likelihood and non-finite derivatives after the labelled probability underflows to zero.

Change

Use log_softmax for labelled class log probabilities. Keep marginal mixtures in log space with logsumexp minus the log sample count, and retain the existing joint-sample aggregation. Preserve the intended averaging of probabilities rather than averaging log probabilities. The probability-input categorical_log_likelihood helper and accuracy/calibration calculations are unchanged.

The operation follows the JAX log_softmax definition.

Validation

• 77 metric tests passed, including 39 new cases. Twenty-five new cases fail on unchanged source; fourteen controls pass before and after.
• Tests cover all three likelihood metrics, eager/JIT use, finite gradients, one/multiple ensemble samples, unequal mixture components, true zero probabilities, and the unchanged probability-input helper.
• Scores are compared with independent float64 SciPy references; labelled-logit derivative sums are checked analytically.
• Scoped Ruff correctness checks and git diff --check pass.
• Run python -m pytest enn/metrics.

Python 3.12.14, JAX/jaxlib 0.11.2, NumPy 2.5.3 and SciPy 1.18.1 on macOS arm64 CPU. The metrics package imports TensorFlow Probability for calibration, so local validation used tfp-nightly 0.26.0.dev20261006; one existing JAX deprecation warning was emitted. No repository dependency files changed. Full ENN training/network tests, static typing, model checkpoints and GPU execution were not run.

Reproduced against c30ae1630bbb41f2f89d061215d851562bc0003c.

[sylvesterkaczmarek]

Following up on this. Is there anything you’d like changed before it can move forward? Happy to address review feedback.