flax
Clean up axis hooks in `nnx.Variable`
#4189
Merged

Loading