flax

5c7fb4c9 - Fix pointless docstring example of nn.checkpoint / nn.remat.

Fix pointless docstring example of nn.checkpoint / nn.remat. It makes no sense to wrap individual matmuls in a remat, you need to wrap the block that uses them.