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.51%. Comparing base (54ba452) to head (98f0ded).

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.52%   84.51%   -0.02%     
==========================================
  Files         178      178              
  Lines       15048    15047       -1     
==========================================
- Hits        12720    12717       -3     
- Misses       2328     2330       +2     

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