No convergence with new commit of ottjax in LRGW solvers unless inner_iterations=10 is set

[selmanozleyen]

I am working on this and I am solving problems with the new ottjax version. You can also see the tests that fail from the CI https://github.com/theislab/moscot/actions/runs/8361086332/job/22888306575#step:5:1934.

The code that fails would look something like this:

ap = (
    AlignmentProblem(adata=adata_space_rotate)
    .prepare(batch_key="batch")
    .solve(epsilon=epsilon, alpha=alpha, rank=rank, initializer=initializer)
)

for prob_key in ap:
    assert ap[prob_key].solution.rank == rank
    assert ap[prob_key].solution.converged

but I noticed if I set inner_iterations=10 it converges

ap = (
    AlignmentProblem(adata=adata_space_rotate)
    .prepare(batch_key="batch")
    .solve(epsilon=epsilon, alpha=alpha, rank=rank, initializer=initializer, inner_iterations=10)
)

for prob_key in ap:
    assert ap[prob_key].solution.rank == rank
    assert ap[prob_key].solution.converged

The new versions of LRGW solvers in ottjax handle inner_iterations differently. This is one example test that fails

moscot/tests/problems/space/test_alignment_problem.py

Line 81 in a5187c0

def test_solve_balanced(

[giovp]

hi @selmanozleyen , yes, you probably noticed this ott-jax/ott@036add2 , we might have to update the default to mirror the ott-jax default

[selmanozleyen]

I think it mirrors the default but the problem is it doesn't converge when it's a higher number like 2000. It does when it is 10. Idk if this is expected

[selmanozleyen]

Ok, I will set the inner iterations to 10 as in the tests of ott-jax. Defaults are already passing.