Serialization support

[cgarciae]

• Adds optional flax serialization support (only if the dependency is available)
• Prepares to add key_path support once jax releases this API.

[codecov]

Codecov Report

Merging #2 (adfdb6e) into main (29b7f43) will increase coverage by 0.26%.
The diff coverage is 97.50%.

@@            Coverage Diff             @@
##             main       #2      +/-   ##
==========================================
+ Coverage   95.65%   95.91%   +0.26%     
==========================================
  Files           2        2              
  Lines          69       98      +29     
==========================================
+ Hits           66       94      +28     
- Misses          3        4       +1     

Impacted Files	Coverage Δ
simple_pytree/pytree.py	95.87% <97.50%> (+0.28%)	⬆️

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more