Maturing stax

[alexbw]

Meta-issue for tracking potential improvements to stax (feel free to rename however you like).

• Metrics Metrics calculations in stax #133
• Model serialization Serializing stax models #134
• Logging Use tracing for program logging and benchmarking #136
• TensorBoard integration Integration with TensorBoard #135
• Pre-trained models Loading pre-trained ONNX models #138
• BatchNorm layer has incorrect test-time behavior Batch norm layer in stax does not compute moving averages for test-time #139

[mattjj]

We should probably discuss whether these features fit into stax. It makes sense to have tools and examples in the ecosystem, but we don’t want to build yet another monolith. Of these, only the batch norm issue and maybe (but not certainly) serialization seem like they are related to neural network building.

[hawkinsp]

We're not planning to work on the current iteration of stax (in favor of supporting other libraries such as Flax and Haiku), so I think we can close this issue.