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[bot]]

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 85.21%. Comparing base (2cdae69) to head (ebe6f70).

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #2868 +/- ## ========================================== - Coverage 85.24% 85.21% -0.03% ========================================== Files 166 166 Lines 14229 14227 -2 ========================================== - Hits 12129 12124 -5 - Misses 2100 2103 +3 ``` | [Files](https://app.codecov.io/gh/scverse/scvi-tools/pull/2868?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) | Coverage Δ | | |---|---|---| | [src/scvi/module/\_vae.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/2868?src=pr&el=tree&filepath=src%2Fscvi%2Fmodule%2F_vae.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvbW9kdWxlL192YWUucHk=) | `95.56% <100.00%> (-0.04%)` | :arrow_down: | ... and [1 file with indirect coverage changes](https://app.codecov.io/gh/scverse/scvi-tools/pull/2868/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse)