flax

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

#5179

Open