feat(module): change base module to support jax device removal

[martinkim0]

jax 0.4.27 removed the device method of arrays in favor of devices, in order to add support for distributed cpu. this method has been deprecated since jax 0.4.21.

the strategy here may need to be adjusted for multi device settings as we currently just return the first device.

[codecov]

Codecov Report

Attention: Patch coverage is 75.00000% with 1 lines in your changes are missing coverage. Please review.

Project coverage is 88.87%. Comparing base (3087178) to head (ee8dfc1).
Report is 4 commits behind head on main.

❗ Current head ee8dfc1 differs from pull request most recent head 877483d. Consider uploading reports for the commit 877483d to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2787      +/-   ##
==========================================
- Coverage   88.90%   88.87%   -0.03%     
==========================================
  Files         161      161              
  Lines       13339    13342       +3     
==========================================
- Hits        11859    11858       -1     
- Misses       1480     1484       +4     

Files	Coverage Δ
src/scvi/module/base/_base_module.py	93.91% <75.00%> (-0.32%)	⬇️

... and 1 file with indirect coverage changes