flax

[nnx] add some optimizations to graph.py

#4377

Merged