Hungarian Algorithm Infinite Loop with NaN Input

[m-sdr]

The Hungarian algorithm implementation in optax.assignment.hungarian_algorithm enters an infinite loop when the input cost matrix contains only NaN (Not a Number) values.

Reproducible Example:

import optax
import jax.numpy as jnp

cost_matrix_nan = jnp.nan * jnp.array([
    [1, 2, 4.],
    [4, 5, 6],
    [7, 8, 9]
])

row_indices, col_indices = optax.assignment.hungarian_algorithm(cost_matrix_nan)

[mblondel]

To be checked if #1140 fixes this issue or not. @carlosgmartin

[carlosgmartin]

It's not clear to me that this is a bug, since (to my mind) it does not really make sense to ask what's the solution to an assignment problem which consists of or contains nan. What would be the correct output?

(One can always apply nan_to_num to the input matrix to get rid of its nans before passing it to hungarian_algorithm, since the output won't matter anyway.)

[mblondel]

But I think we should not allow the algorithm to enter an infinite loop. The ideal would have been to do error reporting, by returning a status variable (succeeded or failed) but this changes the function signature.

What would be the correct output?

Maybe the identity permutation?

An alternative would be to add the nan_to_num call inside hungarian_algorithm instead.

[carlosgmartin]

Actually, it looks like the newer version of the algorithm in #1140 doesn't loop forever on nans. So this should be fixed after that is merged.