feat(nnx): add preferred_element_type to attention.py and recurrent.py, out_sharding to recurrent.py#5179

aarushisingh04 wants to merge 4 commits intogoogle:maingoogle/flax:mainfrom

aarushisingh04:nnx-technical-supportaarushisingh04/flax:nnx-technical-supportCopy head branch name to clipboard