flax

[nnx] improved graph update mechanism

#3759

Merged