Check that all gradient methods run without error.
(grad)
| 113 | @pytest.mark.skipif(not torch, reason="torch not installed") |
| 114 | @pytest.mark.parametrize("grad", ["detach", "envelope", "autodiff", "last_step"]) |
| 115 | def test_gradients_torch(grad): |
| 116 | """Check that all gradient methods run without error.""" |
| 117 | batchsize = 2 |
| 118 | n = 4 |
| 119 | d = 2 |
| 120 | for solver in ["sinkhorn", "log_sinkhorn"]: |
| 121 | X = torch.randn((batchsize, n, d), requires_grad=True) |
| 122 | M = dist_batch(X, X) |
| 123 | res = solve_batch(M, reg=0.1, max_iter=10, tol=1e-5, grad=grad, solver=solver) |
| 124 | loss = res.value_linear.sum() |
| 125 | loss_plan = res.plan.sum() |
| 126 | if grad == "detach": |
| 127 | assert loss.grad == None |
| 128 | elif grad == "envelope": |
| 129 | loss.backward() |
| 130 | assert X.grad is not None |
| 131 | elif grad in ["autodiff", "last_step"]: |
| 132 | loss_plan.backward() |
| 133 | assert X.grad is not None |
| 134 | |
| 135 | |
| 136 | def test_backend(nx): |
nothing calls this directly
no test coverage detected