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.96%   84.51%   -0.46%     
==========================================
  Files         178      178              
  Lines       15048    15047       -1     
==========================================
- Hits        12786    12717      -69     
- Misses       2262     2330      +68     

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

... and 1 file with indirect coverage changes

---- 🚨 Try these New Features:

• Flaky Tests Detection - Detect and resolve failed and flaky tests