Exclusion of observed variables from posterior

[vitkl]

Add tests verifying that observed variables are excluded.

[vitkl]

Hi @adamgayoso @jjhong922

Do you know where the JAX error is coming from and how to solve it? Tests fail at scvi-tools import. Thanks!

[adamgayoso]

can look into it, pinning scvi-tools will not solve this problem. I would leave it >0.15.0

[vitkl]

I see. When did you introduce jax?

[vitkl]

Maybe JAX version needs pinning.

In this test, it is:

Collecting jax>=0.3
  Downloading jax-0.3.7.tar.gz (944 kB)

but in scvi-tools latest successful commit it is:

Collecting jax>=0.3
  Downloading jax-0.3.6.tar.gz (936 kB)

[adamgayoso]

0.15.0 is when we added it, it looks like the issue is optax 0.1.2 released 3 days ago, as our CI is using 0.1.1

[adamgayoso]

oh weird it could also be jax, looks like 0.3.7 is like half-released? It's not released on github but is on pypi

[adamgayoso]

jax-ml/jax#10326

[vitkl]

I see. I will try to fix jax and jaxlib here - and wait for updates.

Did you consider making jax optional for scvi-tools?

[adamgayoso]

We considered making it optional, but feel it will become a larger part of the framework in the future so decided against it. Overtime it will become a more seamless dependency like pytorch I imagine.

[vitkl]

I would hope so too. Looks like fixing version to earlier jax and jaxlib works.

[adamgayoso]

I mean, stuff like this will get sorted out on its own in a few days. A few years ago we all used to have problems with numba, which seems to have been somewhat smoothed out.

[vitkl]

I see. Thanks @adamgayoso ! I will remove fixed versions before merging this.

[vitkl]

WRT exclusion of observed variables:

This appears to be not an issue for standard cell2location - but looks like it is a problem for Messenger AutoGuides: there is no filtering of observed variables in the changes introduces to posterior sampling is this PR scverse/scvi-tools#1267. So filtering needs to be reintroduced (if possible) or that PR reverted.