flax

ea56e5e7 - Add nnx_graph_mode config flag

Add nnx_graph_mode config flag Add a global config flag to control whether NNX APIs default to graph-mode or tree-mode. All public APIs with a `graph` parameter now accept `graph: bool | None = None` and resolve from the config when unset. PiperOrigin-RevId: 872982185