Metrics calculations in stax

[alexbw]

The delta between stax and keras is lower than I originally expected. A functionality stax does not currently have is the calculation of metrics, such as accuracy or precision. See the limited functionality provided by Keras: https://keras.io/metrics/

[mattjj]

If I understand correctly, I don’t think this is within the scope of stax. Stax is just for building neural network functions. Applying those functions for the purpose of training, eval, or anything else doesn’t seem like it needs a separate feature (Python is already good at applying functions). Some of our examples have evaluation.

I’ll close this issue but reopen it if I’m wrong.