flax

remove aliases from nnx.graph in favor of nnx.compat

#5348

Merged