flax

3d634412 - Add tree-mode support for nnx.vmap

Add tree-mode support for nnx.vmap Adds a `graph: bool` parameter to `nnx.vmap`. When `graph=False`, vmap bypasses the graph protocol and operates directly on pytrees, similar to the existing tree-mode support for jit and remat. PiperOrigin-RevId: 870202536