Negative binomial regression on numerical variables

Hi all,

I am trying to use scVI to do negative binomial regression on numerical variables as well as correcting sample batches. Say,

Full model: ~ pseudotime + sampleID
Reduced model: ~ sampleID

Here, the pseudotime is a numerical variable that has been estimated for single cells, and the sampleID is to correct the sample batches. The goal here is to detect pseudotime-correlated genes by excluding sample batches. I understand scVI implements an interface differential_expression() to do the differential gene expression after batch correction, but it is used for categorical groups if I understand it correctly. Could you please help suggest how could I implement a negative binomial regression on continuous numerical variables using scVI? Thanks.