feat: continuous, embedded covariates re-injected during training

[ori-kron-wis]

Added other covariates types (continuous, embedded) to be able to reinjected, and not just to the input layer

[ori-kron-wis]

close https://github.com/scverse/scvi-tools/issues/3021

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 97.82609% with 2 lines in your changes missing coverage. Please review.

Project coverage is 84.35%. Comparing base (795297e) to head (77df3e2).

Files with missing lines	Patch %	Lines
src/scvi/nn/_base_components.py	95.00%	2 Missing :warning:

Additional details and impacted files

```diff @@ Coverage Diff @@ ## main #3032 +/- ## ========================================== - Coverage 84.80% 84.35% -0.46% ========================================== Files 173 173 Lines 14797 14806 +9 ========================================== - Hits 12549 12489 -60 - Misses 2248 2317 +69 ``` | [Files with missing lines](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?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/external/contrastivevi/\_module.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fexternal%2Fcontrastivevi%2F_module.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvZXh0ZXJuYWwvY29udHJhc3RpdmV2aS9fbW9kdWxlLnB5) | `98.61% <100.00%> (ø)` | | | [src/scvi/external/methylvi/\_base\_components.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fexternal%2Fmethylvi%2F_base_components.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvZXh0ZXJuYWwvbWV0aHlsdmkvX2Jhc2VfY29tcG9uZW50cy5weQ==) | `100.00% <100.00%> (ø)` | | | [src/scvi/external/methylvi/\_module.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fexternal%2Fmethylvi%2F_module.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvZXh0ZXJuYWwvbWV0aHlsdmkvX21vZHVsZS5weQ==) | `81.08% <100.00%> (ø)` | | | [src/scvi/module/\_mrdeconv.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fmodule%2F_mrdeconv.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvbW9kdWxlL19tcmRlY29udi5weQ==) | `95.13% <100.00%> (ø)` | | | [src/scvi/module/\_multivae.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fmodule%2F_multivae.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvbW9kdWxlL19tdWx0aXZhZS5weQ==) | `82.22% <100.00%> (ø)` | | | [src/scvi/module/\_peakvae.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fmodule%2F_peakvae.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvbW9kdWxlL19wZWFrdmFlLnB5) | `96.22% <100.00%> (+4.71%)` | :arrow_up: | | [src/scvi/module/\_scanvae.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fmodule%2F_scanvae.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvbW9kdWxlL19zY2FudmFlLnB5) | `86.23% <100.00%> (+0.84%)` | :arrow_up: | | [src/scvi/module/\_totalvae.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fmodule%2F_totalvae.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvbW9kdWxlL190b3RhbHZhZS5weQ==) | `87.54% <100.00%> (ø)` | | | [src/scvi/module/\_vae.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?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=) | `94.50% <100.00%> (+0.02%)` | :arrow_up: | | [src/scvi/module/\_vaec.py](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree&filepath=src%2Fscvi%2Fmodule%2F_vaec.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NjdmkvbW9kdWxlL192YWVjLnB5) | `85.05% <100.00%> (-0.34%)` | :arrow_down: | | ... and [1 more](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) | | ... and [5 files with indirect coverage changes](https://app.codecov.io/gh/scverse/scvi-tools/pull/3032/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse)

[ori-kron-wis]

This kind of change requires touching the module for each model. I don't think there is anywhere around it. Its risky and this is why I implemented it only for scvi & scanvi at the moment (maybe others will work also but I didn't test thoroughly). the rest of the changed modules are mainly things in function headers and placeholders to able to cope with the major changes in base_components (i.e cont_cov is usually None and cant be injected to their hidden layers).

To summarise we have 2 parameters to think of:

• deeply_inject (bool): covariates (continuous or categorial whether being encoded or not) are always injected to the first layer. If this parameter is True they will also be injected to the hidden layer of the encoder/decoder.
• encode_covariates (bool): this is only relevant for the categorial covariates. If this parameter is false, they will be injected as one-hot, if true, they will be encoded as an embedded matrix (this is dependent on other meta-parameters such as the size of embedded vector and the batch_representation which should be "embedded" as well if we want it, otherwise it will remain one-hot)

@Hrovatin please try to use this branch

[Hrovatin]

@ori-kron-wis I think there are things that need to be fixed as per my review comments. To summarise: I think the behaviour should be as follows:

• Module passes to layers expression, continuous covariates (can be None), and categorical covariates that need to be one-hot encoded (can be empty list)
• Module is the one that decides if batch will be embedded and concatenated to continuous covariates or added to the categorical list
• Layers should not check how the batch is encoded as their init params already tell the number of continous/categorical covariates as determined by the Module.
• The layers then either inject covariates (continuous and one-hote encoded categorical) into all layers or only add to the first one. - So Module should not be concatenating continuous covariates to expression in advance.

@canergen can you confirm?