use jax.Array type for rng keys

[chiamp]

for more details, see the Jax PR and JEP 9263

[codecov-commenter]

Codecov Report

Merging #3354 (dbc9254) into main (ae402f3) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main    #3354   +/-   ##
=======================================
  Coverage   83.03%   83.03%           
=======================================
  Files          55       55           
  Lines        6521     6521           
=======================================
  Hits         5415     5415           
  Misses       1106     1106           

Files Changed	Coverage Δ
flax/linen/module.py	92.02% <100.00%> (ø)
flax/linen/recurrent.py	97.30% <100.00%> (ø)
flax/linen/stochastic.py	96.77% <100.00%> (ø)

[IvyZX]

Thanks for making the change!