flax

cfdf1d41 - [nnx] add unsafe_graph_cache