Improve Typing support

[cgarciae]

What does this PR do?

• Makes AxisMetadata generic over the unbox type and improves some method signatures.
• Adds overload signatures for Module.param, Module.variable, Module.apply, Scope.param, and Scope.variable.

[levskaya]

Looks good - but I think we need to do the same thing for scope.variable and self.variable ?

[codecov-commenter]

Codecov Report

Merging #3242 (36e5141) into main (675e34d) will decrease coverage by 0.14%.
Report is 3 commits behind head on main.
The diff coverage is 73.97%.

@@            Coverage Diff             @@
##             main    #3242      +/-   ##
==========================================
- Coverage   82.35%   82.21%   -0.14%     
==========================================
  Files          54       54              
  Lines        6087     6147      +60     
==========================================
+ Hits         5013     5054      +41     
- Misses       1074     1093      +19     

Files Changed	Coverage Δ
flax/linen/module.py	91.33% <67.64%> (-1.13%)	⬇️
flax/core/scope.py	88.91% <68.00%> (-1.14%)	⬇️
flax/core/meta.py	95.40% <100.00%> (+0.10%)	⬆️
flax/linen/attention.py	96.06% <100.00%> (ø)