flax

4870afe5 - Fixed nnx.remat docstring rendering