Have _graph_flatten respect nnx.data declarations#5136

samanklesaria wants to merge 2 commits intogoogle:maingoogle/flax:mainfrom

samanklesaria:issues/5128samanklesaria/flax:issues/5128Copy head branch name to clipboard