flax

Modify NNX to use id(variable) instead of nnx.Variables as dictionary

#4814

Merged