flax

befe5718 - Merge pull request #5008 from mohsinm-dev:fix-nnx-tabulate-variable-hooks