| 235 | |
| 236 | @pytest.mark.skipif(not torch, reason="torch no installed") |
| 237 | def test_solve_envelope(): |
| 238 | n_samples_s = 10 |
| 239 | n_samples_t = 7 |
| 240 | n_features = 2 |
| 241 | rng = np.random.RandomState(0) |
| 242 | |
| 243 | x = rng.randn(n_samples_s, n_features) |
| 244 | y = rng.randn(n_samples_t, n_features) |
| 245 | a = ot.utils.unif(n_samples_s) |
| 246 | b = ot.utils.unif(n_samples_t) |
| 247 | M = ot.dist(x, y) |
| 248 | |
| 249 | a = torch.tensor(a, requires_grad=True) |
| 250 | b = torch.tensor(b, requires_grad=True) |
| 251 | M = torch.tensor(M, requires_grad=True) |
| 252 | |
| 253 | sol0 = ot.solve(M, a, b, reg=10, grad="envelope") |
| 254 | sol0.value.backward() |
| 255 | |
| 256 | gM0 = M.grad.clone() |
| 257 | ga0 = a.grad.clone() |
| 258 | gb0 = b.grad.clone() |
| 259 | |
| 260 | a = torch.tensor(a, requires_grad=True) |
| 261 | b = torch.tensor(b, requires_grad=True) |
| 262 | M = torch.tensor(M, requires_grad=True) |
| 263 | |
| 264 | sol = ot.solve(M, a, b, reg=10, grad="autodiff") |
| 265 | sol.value.backward() |
| 266 | |
| 267 | gM = M.grad.clone() |
| 268 | ga = a.grad.clone() |
| 269 | gb = b.grad.clone() |
| 270 | |
| 271 | # Note, gradients are invariant to change in constant so we center them |
| 272 | assert torch.allclose(gM0, gM) |
| 273 | assert torch.allclose(ga0 - ga0.mean(), ga - ga.mean()) |
| 274 | assert torch.allclose(gb0 - gb0.mean(), gb - gb.mean()) |
| 275 | |
| 276 | |
| 277 | @pytest.mark.parametrize( |