flax

c8cf3f50 - [nnx] add ref_cache to graph_unflatten