| 150 | |
| 151 | @pytest.mark.skipif(not torch, reason="torch no installed") |
| 152 | def test_solve_last_step(): |
| 153 | n_samples_s = 10 |
| 154 | n_samples_t = 7 |
| 155 | n_features = 2 |
| 156 | rng = np.random.RandomState(0) |
| 157 | |
| 158 | x = rng.randn(n_samples_s, n_features) |
| 159 | y = rng.randn(n_samples_t, n_features) |
| 160 | a = ot.utils.unif(n_samples_s) |
| 161 | b = ot.utils.unif(n_samples_t) |
| 162 | M = ot.dist(x, y) |
| 163 | |
| 164 | # Check that last_step and autodiff give the same result and similar gradients |
| 165 | a = torch.tensor(a, requires_grad=True) |
| 166 | b = torch.tensor(b, requires_grad=True) |
| 167 | M = torch.tensor(M, requires_grad=True) |
| 168 | |
| 169 | sol0 = ot.solve(M, a, b, reg=10, grad="autodiff") |
| 170 | sol0.value.backward() |
| 171 | |
| 172 | gM0 = M.grad.clone() |
| 173 | ga0 = a.grad.clone() |
| 174 | gb0 = b.grad.clone() |
| 175 | |
| 176 | a = torch.tensor(a, requires_grad=True) |
| 177 | b = torch.tensor(b, requires_grad=True) |
| 178 | M = torch.tensor(M, requires_grad=True) |
| 179 | |
| 180 | sol = ot.solve(M, a, b, reg=10, grad="last_step") |
| 181 | sol.value.backward() |
| 182 | |
| 183 | gM = M.grad.clone() |
| 184 | ga = a.grad.clone() |
| 185 | gb = b.grad.clone() |
| 186 | |
| 187 | # Note, gradients are invariant to change in constant so we center them |
| 188 | cos = torch.nn.CosineSimilarity(dim=0, eps=1e-6) |
| 189 | tolerance = 0.96 |
| 190 | assert cos(gM0.flatten(), gM.flatten()) > tolerance |
| 191 | assert cos(ga0 - ga0.mean(), ga - ga.mean()) > tolerance |
| 192 | assert cos(gb0 - gb0.mean(), gb - gb.mean()) > tolerance |
| 193 | |
| 194 | assert torch.allclose(sol0.plan, sol.plan) |
| 195 | assert torch.allclose(sol0.value, sol.value) |
| 196 | assert torch.allclose(sol0.value_linear, sol.value_linear) |
| 197 | assert torch.allclose(sol0.potentials[0], sol.potentials[0]) |
| 198 | assert torch.allclose(sol0.potentials[1], sol.potentials[1]) |
| 199 | |
| 200 | with pytest.raises(ValueError): |
| 201 | ot.solve(M, a, b, grad="last_step", max_iter=0, reg=10) |
| 202 | |
| 203 | |
| 204 | @pytest.mark.skipif(not torch, reason="torch no installed") |