Linen+ Extensions for the Flax linen library. Install pip install git+https://github.com/ikostrikov/linenplus Examples You can find a bunch of examples in the examples folder.