Grid geometry handling of 0 weights seems buggy

[marcocuturi]

Grid geometry bugs with 0 weights

Running this test with some values set to 0 results in nan but it shouldn't. This test is a simplified version of the test insinkhorn_test.py, setting two values to 0.

ns = [6, 7, 11]
xs = [
    jax.random.normal(jax.random.PRNGKey(i), (n,))
    for i, n in enumerate(ns)
]
cost_fn=None
geom = ott.geometry.grid.Grid(xs, epsilon=0.1)
a = jax.random.uniform(jax.random.PRNGKey(0), (geom.shape[0],))
b = jax.random.uniform(jax.random.PRNGKey(1), (geom.shape[0],))
a = a.at[10].set(0.0)
b = a.at[13].set(0.0)
a, b = a / jnp.sum(a), b / jnp.sum(b)
lin_prob = ott.problems.linear.linear_problem.LinearProblem(geom, a=a, b=b)
solver = ott.solvers.linear.sinkhorn.Sinkhorn()
out = solver(lin_prob)

# Recover full cost matrix by applying it to columns of identity matrix.
cost_matrix = geom.apply_cost(jnp.eye(geom.shape[0]))
# Recover full transport by applying it to columns of identity matrix.
transport_matrix = out.apply(jnp.eye(geom.shape[0]))
cost = jnp.sum(transport_matrix * cost_matrix)

latter is nan.

[michalk8]

I briefly looked into this, looks like numerical imprecision, on float64 it seems to be ok;
the out.errors also seem to confirm this (seems stable for a while, then it shoots up).

[marcocuturi]

Thanks for checking as well! Yes, I agree, this is also what a closer look was showing... The only thing that still worries me a bit is that even for large epsilon there still seems to be some issues! Will take a closer look.

[marcocuturi]

Thanks. Motivation to recenter potentials can be found in this PR, where instantiating the scalings from very large potential values was causing overflow, and therefore nan cost values.