flax

26d5e68d - [nnx] fix Accuracy on eager mode