[MRG] Fix GMM rand map overflow

[jonathan-legrand]

Types of changes

I simplified the computation of the $T_{rand}$ map and used the logsumexp trick to avoid overflows. I also added a test case that triggers overflow errors with the current main branch implementation and passes with the changes introduced in this request.

Motivation and context / Related issue

I ran into overflow errors when transporting gaussian mixtures with the $T_{rand}$ map. They occur when a point of the source domain $x$ is far away from one of the target components.
The current implementation computes: log_diff = log_g[:, None] - log_g[None, :] which is effectively:

$$ \log( g_i(x) ) - \log(g_j(x)) = \log( \frac{g_i(x)} {g_j(x) }) $$

and then exponentiates this quantity : weighted_exp = w_s[:, None] * nx.exp(log_diff)

When the ratio $g_i(x)/g_j(x)$ is very large, the exp overflows.

I rewrote this computation using the logsumexp trick. I saw that a logsumexp function is available through nx but it does not accept weights, which are required for computing the $T_{rand}$ attribution probability, so I implemented one which is backend agnostic.

How has this been tested (if it applies)

The logsumexp has been tested for high (710) and low (0) logits and it behaves as expected. The tests pass for numpy and torch backends. I had to skip jax tests because the functions of the gmm module use a lot of array assignments.

PR checklist

• I have read the CONTRIBUTING document.
• The documentation is up-to-date with the changes I made (check build artifacts).
• All tests passed, and additional code has been covered with new tests.
• I have added the PR and Issue fix to the RELEASES.md file.