fixed NNX decode and dynamic slicing

[chiamp]

Added @giovannic's dynamic slicing fix to NNX.

Fixed NNX init_cache since it was initializing a cache with trailing dimensions (..., self.num_heads, self.out_features) instead of (..., self.num_heads, self.head_dim).

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (008c3cc) 56.32% compared to head (f17b4fe) 56.70%.
Report is 22 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3576      +/-   ##
==========================================
+ Coverage   56.32%   56.70%   +0.38%     
==========================================
  Files         100      100              
  Lines       11951    12006      +55     
==========================================
+ Hits         6731     6808      +77     
+ Misses       5220     5198      -22     

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