Remove jnp.DeviceArray since it is deprecated. (#24875)
* Remove jnp.DeviceArray since it is deprecated.
* Replace all instances of jnp.DeviceArray with jax.Array
* Update src/transformers/models/bert/modeling_flax_bert.py
---------
Co-authored-by: Sanchit Gandhi <93869735+sanchit-gandhi@users.noreply.github.com>
