[nnx] always fold_in on fork + new ForkedKeys return type

[cgarciae]

What does this PR do?

• Rngs.fork now returns a single ForkedKeys type instead of different types depending on the overload (overloads are removed). ForkedKeys is a Mapping[str, Array] with the keys, and contains .splits and .broadcasts attributes of type dict[str, Array] for cases when you need to tell apart the splits from the broadcasts.
• Rngs.fork now folds all keys so there is no need to maintain / propagate hashable data.
• Because of the previous RngStream.counts: list[int] is replaced with RngStream.count: int.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 94.73684% with 5 lines in your changes are missing coverage. Please review.

Project coverage is 59.05%. Comparing base (1abfa87) to head (3e69350).
Report is 2 commits behind head on main.

Files	Patch %	Lines
flax/experimental/nnx/nnx/rnglib.py	94.00%	3 Missing ⚠️
flax/experimental/nnx/nnx/transforms.py	75.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3722      +/-   ##
==========================================
+ Coverage   59.00%   59.05%   +0.04%     
==========================================
  Files         103      103              
  Lines       12422    12438      +16     
==========================================
+ Hits         7330     7345      +15     
- Misses       5092     5093       +1     

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