Add static_argnums to nn.checkpoint

[cgarciae]

What does this PR do?

Fixes #2452. Adds a static_argnums to nn.checkpoint and pipes it to jax.remat, copies argument docstring from JAX.

[cgarciae]

I want to add another test, lets wait before merging.

[cgarciae]

I discussed this PR with @jheek, final behaviour:

• subtracts -1 from each static_argnum in transforms.checkpoint because self (the Module) is never passed to the lifted transformation.
• adds +2 to each static_argnum in lift.checkpoint because 2 arguments (variable_groups, rng_groups) are always passed to the inner function in the lifted transformation.

Added an additional test. Should be ready for a final review + internal merge.

[codecov-commenter]

Codecov Report

Merging #2457 (f384b33) into main (e320e11) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main    #2457   +/-   ##
=======================================
  Coverage   79.66%   79.67%           
=======================================
  Files          49       49           
  Lines        4982     4984    +2     
=======================================
+ Hits         3969     3971    +2     
  Misses       1013     1013           

Impacted Files	Coverage Δ
flax/core/lift.py	95.81% <100.00%> (+<0.01%)	⬆️
flax/linen/transforms.py	94.06% <100.00%> (+0.02%)	⬆️

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.