| 189 | |
| 190 | |
| 191 | def test_entropic_partial_wasserstein_gradient(): |
| 192 | if torch: |
| 193 | n_samples = 40 |
| 194 | |
| 195 | mu = np.array([0, 0]) |
| 196 | cov = np.array([[1, 0], [0, 2]]) |
| 197 | |
| 198 | xs = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov) |
| 199 | xt = ot.datasets.make_2D_samples_gauss(n_samples, mu, cov) |
| 200 | |
| 201 | M = torch.tensor(ot.dist(xs, xt), requires_grad=True, dtype=torch.float64) |
| 202 | |
| 203 | p = torch.tensor(ot.unif(n_samples), requires_grad=True, dtype=torch.float64) |
| 204 | q = torch.tensor(ot.unif(n_samples), requires_grad=True, dtype=torch.float64) |
| 205 | |
| 206 | m = 0.5 |
| 207 | reg = 1 |
| 208 | |
| 209 | _, log = ot.partial.entropic_partial_wasserstein( |
| 210 | p, q, M, m=m, reg=reg, log=True |
| 211 | ) |
| 212 | |
| 213 | log["partial_w_dist"].backward() |
| 214 | |
| 215 | assert M.grad is not None |
| 216 | assert p.grad is not None |
| 217 | assert q.grad is not None |
| 218 | assert M.grad.shape == M.shape |
| 219 | assert p.grad.shape == p.shape |
| 220 | assert q.grad.shape == q.shape |
| 221 | |
| 222 | |
| 223 | def test_partial_gromov_wasserstein(): |