Many improvements to Linen, and the old flax.nn is officially reprecated!
Notably, there's a clean API for extracting intermediates from modules
defined using @nn.compact, a more ergonomic API for using Batch Norm and Dropout in modules
defined using setup, support for MultiOptimizer with Linen, and multiple safety, performance
and error message improvements.
See the CHANGELOG for more details