diff options
| author | Clément Bonet <32179275+clbonet@users.noreply.github.com> | 2023-04-18 18:01:19 +0200 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2023-04-18 18:01:19 +0200 |
| commit | 9aa96c8247afd6e98d8bd470a6adb1be0f1c467e (patch) | |
| tree | 3f213c8d844d6f24f88c83deebec55f45391e4f9 /ot/lp/solver_1d.py | |
| parent | 1078dcc3530a7f95fd77d19d115d46f39c2574bc (diff) | |
[MRG] Fix Bug binary_search_circle on GPU and Gradients (#457)
* W circle + SSW
* Tests + Example SSW_1
* Example Wasserstein Circle + Tests
* Wasserstein on the circle wrt Unif
* Example SSW unif
* pep8
* np.linalg.qr for numpy < 1.22 by batch + add python3.11 to tests
* np qr
* rm test python 3.11
* update names, tests, backend transpose
* Comment error batchs
* semidiscrete_wasserstein2_unif_circle example
* torch permute method instead of torch.permute for previous versions
* update comments and doc
* doc wasserstein circle model as [0,1[
* Added ot.utils.get_coordinate_circle to get coordinates on the circle in turn
* Bug cuda w_circle + gradient ssw
* Bug cuda w_circle + gradient ssw
* backend detach
* Add PR in Releases.md
---------
Co-authored-by: Rémi Flamary <remi.flamary@gmail.com>
Diffstat (limited to 'ot/lp/solver_1d.py')
| -rw-r--r-- | ot/lp/solver_1d.py | 10 |
1 files changed, 5 insertions, 5 deletions
diff --git a/ot/lp/solver_1d.py b/ot/lp/solver_1d.py index bcfc920..840801a 100644 --- a/ot/lp/solver_1d.py +++ b/ot/lp/solver_1d.py @@ -401,7 +401,7 @@ def roll_cols(M, shifts): n_rows, n_cols = M.shape - arange1 = nx.tile(nx.reshape(nx.arange(n_cols), (1, n_cols)), (n_rows, 1)) + arange1 = nx.tile(nx.reshape(nx.arange(n_cols, type_as=shifts), (1, n_cols)), (n_rows, 1)) arange2 = (arange1 - shifts) % n_cols return nx.take_along_axis(M, arange2, 1) @@ -600,7 +600,7 @@ def binary_search_circle(u_values, v_values, u_weights=None, v_weights=None, p=1 using e.g. ot.utils.get_coordinate_circle(x) - The function runs on backend but tensorflow is not supported. + The function runs on backend but tensorflow and jax are not supported. Parameters ---------- @@ -730,7 +730,7 @@ def binary_search_circle(u_values, v_values, u_weights=None, v_weights=None, p=1 tp[((1 - mask) * (dCp >= 0)) > 0] = tc[((1 - mask) * (dCp >= 0)) > 0] tc[((1 - mask) * (1 - done)) > 0] = (tm[((1 - mask) * (1 - done)) > 0] + tp[((1 - mask) * (1 - done)) > 0]) / 2 - w = ot_cost_on_circle(tc, u_values, v_values, u_cdf, v_cdf, p) + w = ot_cost_on_circle(nx.detach(tc), u_values, v_values, u_cdf, v_cdf, p) if log: return w, {"optimal_theta": tc[:, 0]} @@ -743,7 +743,7 @@ def wasserstein1_circle(u_values, v_values, u_weights=None, v_weights=None, requ takes the value modulo 1. If the values are on :math:`S^1\subset\mathbb{R}^2`, first find the coordinates using e.g. the atan2 function. - The function runs on backend but tensorflow is not supported. + The function runs on backend but tensorflow and jax are not supported. .. math:: W_1(u,v) = \int_0^1 |F_u(t)-F_v(t)-LevMed(F_u-F_v)|\ \mathrm{d}t @@ -864,7 +864,7 @@ def wasserstein_circle(u_values, v_values, u_weights=None, v_weights=None, p=1, using e.g. ot.utils.get_coordinate_circle(x) - The function runs on backend but tensorflow is not supported. + The function runs on backend but tensorflow and jax are not supported. Parameters ---------- |