(feat): Linear GENOT

[ilan-gold]

We add two main new classes here:

1. GENOTLinSolver wraps the underlying OTT-Jax solver in a moscot-usable way
2. GENOTLinProblem at the next level of abstraction up, allowing users to work with AnnData objects directly while also using other parts of the moscot API like policies

To run this PR, simply do pip install -e . as it points directly at the corresponding OTT-JAX branch.

TODOs:

• [x] Finalize status with OTT-Jax PR
• [x] Remove more dead code as this was based on the ICNN implementation (if there is any)
• [x] Maybe add more tests for certain things that I did not cover? I will need to look this over
• [x] Maybe a small notebook?
• [x] Figure out how to squash all these commits. It should be clear where I picked up and the ICNN branch stopped but in any case...would be nice on the eyes. And we definitely want to squash this PR when merged

[ilan-gold]

Some todos:

1. RNG for neural class
2. Validation loop/feature
3. Reasonable defaults

[ilan-gold]

https://github.com/theislab/moscot/actions/runs/9221997646/job/25372233453?pr=662 complains that it can't find diffrax with the minimum version from ott-jax[neural] on python 3.8, and I'm not sure what to make of that since ott-jax tests against 3.8. That being said, diffrax at the minimum version does not support 3.8 it seems: https://github.com/patrick-kidger/diffrax/blob/v0.4.1/setup.py

Thoughts on this from the Kleins? @MUCDK @michalk8

I may be missing something from the testing setup for ott-jax, like it doesn't run neural against 3.8. But figured it would be quicker to ask.

[MUCDK]

https://github.com/theislab/moscot/actions/runs/9221997646/job/25372233453?pr=662 complains that it can't find diffrax with the minimum version from ott-jax[neural] on python 3.8, and I'm not sure what to make of that since ott-jax tests against 3.8. That being said, diffrax at the minimum version does not support 3.8 it seems: https://github.com/patrick-kidger/diffrax/blob/v0.4.1/setup.py

Thoughts on this from the Kleins? @MUCDK @michalk8

I may be missing something from the testing setup for ott-jax, like it doesn't run neural against 3.8. But figured it would be quicker to ask.

what about deprecating python 3.8 , also @giovp ?

[giovp]

yes lets deprecate!

[ilan-gold]

ok @MUCDK !

[selmanozleyen]

@MUCDK do you think GENOT code in ottjax is final? My version in my thesis project is very different from the one in ottjax . If you like maybe we can first discuss possible changes there. Depends on what you want to prioritize.

[MUCDK]

@MUCDK do you think GENOT code in ottjax is final? My version in my thesis project is very different from the one in ottjax . If you like maybe we can first discuss possible changes there. Depends on what you want to prioritize.

Yeah it's not super final but I wouuld prefer to have the merges more incremental as we have been workin gon this PR for almost 2 years.

Let's talk next week in person.

[MUCDK]

Greatl, let's wait for @ilan-gold 's response, and then merge.