()
| 478 | |
| 479 | |
| 480 | def test_sliced_sphere_gradient(): |
| 481 | if torch: |
| 482 | import torch.nn.functional as F |
| 483 | |
| 484 | X0 = torch.randn((20, 3)) |
| 485 | X0 = F.normalize(X0, p=2, dim=-1) |
| 486 | X0.requires_grad_(True) |
| 487 | |
| 488 | X1 = torch.randn((20, 3)) |
| 489 | X1 = F.normalize(X1, p=2, dim=-1) |
| 490 | |
| 491 | sw = ot.sliced_wasserstein_sphere(X1, X0, n_projections=100, p=2) |
| 492 | grad_x0 = torch.autograd.grad(sw, X0)[0] |
| 493 | |
| 494 | assert not torch.any(torch.isnan(grad_x0)) |
| 495 | |
| 496 | |
| 497 | def test_sliced_sphere_unif_values_on_the_sphere(): |