re-factored features arg

[chiamp]

Standardized feature args to in_features and out_features, since NNX uses both in/out_features and features_in/out, while Linen only uses in/out_features.

[codecov-commenter]

Codecov Report

Attention: 2 lines in your changes are missing coverage. Please review.

Comparison is base (2e51a4e) 56.17% compared to head (aa1703b) 56.17%.
Report is 4 commits behind head on main.

Files	Patch %	Lines
flax/experimental/nnx/nnx/nn/attention.py	83.33%	1 Missing ⚠️
flax/linen/attention.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3554   +/-   ##
=======================================
  Coverage   56.17%   56.17%           
=======================================
  Files         100      100           
  Lines       11890    11890           
=======================================
  Hits         6679     6679           
  Misses       5211     5211           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[cgarciae]

This is great, thanks!