| 203 | |
| 204 | @pytest.mark.skipif(not torch, reason="torch no installed") |
| 205 | def test_solve_detach(): |
| 206 | n_samples_s = 10 |
| 207 | n_samples_t = 7 |
| 208 | n_features = 2 |
| 209 | rng = np.random.RandomState(0) |
| 210 | |
| 211 | x = rng.randn(n_samples_s, n_features) |
| 212 | y = rng.randn(n_samples_t, n_features) |
| 213 | a = ot.utils.unif(n_samples_s) |
| 214 | b = ot.utils.unif(n_samples_t) |
| 215 | M = ot.dist(x, y) |
| 216 | |
| 217 | # Check that last_step and autodiff give the same result and similar gradients |
| 218 | a = torch.tensor(a, requires_grad=True) |
| 219 | b = torch.tensor(b, requires_grad=True) |
| 220 | M = torch.tensor(M, requires_grad=True) |
| 221 | |
| 222 | sol0 = ot.solve(M, a, b, reg=10, grad="detach") |
| 223 | |
| 224 | with pytest.raises(RuntimeError): |
| 225 | sol0.value.backward() |
| 226 | |
| 227 | sol = ot.solve(M, a, b, reg=10, grad="autodiff") |
| 228 | |
| 229 | assert torch.allclose(sol0.plan, sol.plan) |
| 230 | assert torch.allclose(sol0.value, sol.value) |
| 231 | assert torch.allclose(sol0.value_linear, sol.value_linear) |
| 232 | assert torch.allclose(sol0.potentials[0], sol.potentials[0]) |
| 233 | assert torch.allclose(sol0.potentials[1], sol.potentials[1]) |
| 234 | |
| 235 | |
| 236 | @pytest.mark.skipif(not torch, reason="torch no installed") |