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 50.00000% with 4 lines in your changes missing coverage. Please review.

Project coverage is 83.25%. Comparing base (09d82ea) to head (241a377).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/scvi/train/_trainingplans.py	50.00%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2931      +/-   ##
==========================================
- Coverage   83.27%   83.25%   -0.02%     
==========================================
  Files         180      180              
  Lines       15187    15193       +6     
==========================================
+ Hits        12647    12649       +2     
- Misses       2540     2544       +4     

Files with missing lines	Coverage Δ
src/scvi/module/_vae.py	94.92% <ø> (ø)
src/scvi/train/_trainingplans.py	93.81% <50.00%> (-0.79%)	⬇️