flax

f265a5e2 - Merge pull request #4374 from 8bitmp3:add-nnx-eval_shap-docstring