add quick_start

[cgarciae]

Changes

• Adds docs/quick_start.ipynb
• Copies pooling.py from Flax and exposes the pooling funcitons.
• Adds value_and_grad.

[codecov]

Codecov Report

Merging #18 (22bb302) into main (3ae486a) will decrease coverage by 1.53%.
The diff coverage is 34.92%.

@@            Coverage Diff             @@
##             main      #18      +/-   ##
==========================================
- Coverage   78.93%   77.40%   -1.53%     
==========================================
  Files          20       21       +1     
  Lines        1837     1890      +53     
==========================================
+ Hits         1450     1463      +13     
- Misses        387      427      +40     

Impacted Files	Coverage Δ
nnx/nn/pooling.py	20.00% <20.00%> (ø)
nnx/transforms.py	83.05% <50.00%> (-2.67%)	⬇️
nnx/__init__.py	100.00% <100.00%> (ø)
nnx/helpers.py	60.00% <100.00%> (ø)
nnx/module.py	78.34% <100.00%> (ø)

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