flax

Have _graph_flatten respect nnx.data declarations (extension of #5140)

#5159

Merged