| 164 | |
| 165 | |
| 166 | def test_partial_wasserstein2_gradient(): |
| 167 | if torch: |
| 168 | n_samples = 40 |
| 169 | |
| 170 | mu = np.array([0, 0]) |
| 171 | cov = np.array([[1, 0], [0, 2]]) |
| 172 | |
| 173 | xs = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov) |
| 174 | xt = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov) |
| 175 | |
| 176 | M = torch.tensor(ot.dist(xs, xt), requires_grad=True, dtype=torch.float64) |
| 177 | |
| 178 | p = torch.tensor(ot.unif(n_samples), dtype=torch.float64) |
| 179 | q = torch.tensor(ot.unif(n_samples), dtype=torch.float64) |
| 180 | |
| 181 | m = 0.5 |
| 182 | |
| 183 | w, log = ot.partial.partial_wasserstein2(p, q, M, m=m, log=True) |
| 184 | |
| 185 | w.backward() |
| 186 | |
| 187 | assert M.grad is not None |
| 188 | assert M.grad.shape == M.shape |
| 189 | |
| 190 | |
| 191 | def test_entropic_partial_wasserstein_gradient(): |