Add benchmark on state traversal, and a readme

[IvyZX]

Adds a benchmark script showing that the NNX state pytree traversal improved ~6x on flatten and ~3x on tree-map, after the latest JAX release.

Running the script in jaxlib==0.4.35 jax==0.4.35 gives:

### tree_flatten_with_path ###
total time: 3.397632122039795
time per step: 3397.63 µs
time per layer: 849.41 µs
### tree_map_with_path ###
total time: 4.1260809898376465
time per step: 4126.08 µs
time per layer: 1031.52 µs

And running in jaxlib==0.4.36 jax==0.4.37 gives:

### tree_flatten_with_path ###
total time: 0.5094578266143799
time per step: 509.46 µs
time per layer: 127.36 µs
### tree_map_with_path ###
total time: 1.18902587890625
time per step: 1189.03 µs
time per layer: 297.26 µs