style: add initial tensor shape typing to vae

[martinkim0]

closes #2541.

draft implementation of adding jaxtyping tensor shape annotation to vae

[martinkim0]

Pros:

• Same syntax across JAX, PyTorch, and NumPy (not sure about Xarray or Sparse yet)
• Track shape annotations for tensors, which can make it easier to catch bugs earlier

Cons:

• Can be tedious to add typing to every function - probably need to decide when it's recommended and when it's optional
• Can sometimes lead to more verbose code

Other considerations:

• Need to use the same axis names across all the models for consistency, so need a good system to sync these and manage axis name conflicts

[codecov]

Codecov Report

Attention: Patch coverage is 90.00000% with 2 lines in your changes missing coverage. Please review.

Project coverage is 84.29%. Comparing base (38e1d65) to head (e88ae3c).

Files with missing lines	Patch %	Lines
src/scvi/module/_vae.py	90.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2868      +/-   ##
==========================================
- Coverage   84.80%   84.29%   -0.52%     
==========================================
  Files         173      173              
  Lines       14797    14794       -3     
==========================================
- Hits        12549    12470      -79     
- Misses       2248     2324      +76     

Files with missing lines	Coverage Δ
src/scvi/module/_vae.py	93.67% <90.00%> (-0.82%)	⬇️

... and 5 files with indirect coverage changes