feat(train) Add support for torch.compile (EXPERIMENTAL)

[canergen]

No description provided.

[canergen]

@ori-kron-wis Can you add tests for all pytorch models (not pyro and not jax) for compile. Can you check speed improvements on your end? You should execute it with: model.train(accelerator='cuda', plan_kwargs={'n_epochs_kl_warmup': 100, 'compile': True}, datasplitter_kwargs={'drop_last': True})

[canergen]

Needs tests like: model2.train(accelerator='cuda', batch_size=5000, max_epochs=100, train_size=0.9, plan_kwargs={'n_epochs_kl_warmup': 100, 'compile': True}, datasplitter_kwargs={'drop_last': True}) and then get_elbo, get_reconstruction_loss, get_latent.

[ori-kron-wis]

I added torch compile tests for most models (of course not working with the github action due to that error) - on new servers, it worked fine and was faster, although the compile part will add some overhead.

Currently pyro test not working on a multi GPU machine. Need to see why (only test_pyro_bayesian_regression). once we remove it everything works (it should be passed here)

[codecov]

Codecov Report

Attention: Patch coverage is 70.58824% with 5 lines in your changes missing coverage. Please review.

Project coverage is 84.33%. Comparing base (e0e0209) to head (33453ca).

Files with missing lines	Patch %	Lines
src/scvi/train/_trainingplans.py	71.42%	4 Missing ⚠️
src/scvi/external/poissonvi/_model.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2931      +/-   ##
==========================================
- Coverage   84.35%   84.33%   -0.03%     
==========================================
  Files         167      167              
  Lines       14490    14501      +11     
==========================================
+ Hits        12223    12229       +6     
- Misses       2267     2272       +5     

Files with missing lines	Coverage Δ
src/scvi/module/_vae.py	94.48% <ø> (ø)
src/scvi/external/poissonvi/_model.py	71.25% <66.66%> (-0.55%)	⬇️
src/scvi/train/_trainingplans.py	93.84% <71.42%> (-0.75%)	⬇️