Use pyink

[cgarciae]

What does this PR do?

• Adds pyink as a pre-commit hook.
• Adds some pyink configuration in pyproject.toml
• Formats entire project using pyink.

[codecov-commenter]

Codecov Report

Merging #3216 (fe6112a) into main (c8bb930) will decrease coverage by 0.09%.
The diff coverage is 77.92%.

@@            Coverage Diff             @@
##             main    #3216      +/-   ##
==========================================
- Coverage   82.32%   82.24%   -0.09%     
==========================================
  Files          54       54              
  Lines        6071     6076       +5     
==========================================
- Hits         4998     4997       -1     
- Misses       1073     1079       +6     

Impacted Files	Coverage Δ
flax/configurations.py	89.65% <ø> (ø)
flax/core/__init__.py	100.00% <ø> (ø)
flax/core/partial_eval.py	100.00% <ø> (ø)
flax/ids.py	96.00% <ø> (ø)
flax/testing/benchmark.py	0.00% <0.00%> (ø)
flax/training/common_utils.py	0.00% <0.00%> (ø)
flax/training/dynamic_scale.py	0.00% <0.00%> (ø)
flax/training/lr_schedule.py	0.00% <0.00%> (ø)
flax/training/prefetch_iterator.py	0.00% <0.00%> (ø)
flax/training/train_state.py	83.33% <0.00%> (ø)
... and 39 more