flax

8f8ec30a - Add tree-mode support to iter_graph

Add tree-mode support to iter_graph Add a graph: bool | None parameter to `iter_graph` and other downstream functions such as `iter_modules`, `split_rngs`, `fork_rngs`, `backup_keys`, and `reseed`. PiperOrigin-RevId: 873211727