diff options
Diffstat (limited to 'test')
| -rw-r--r-- | test/test_backend.py | 12 | ||||
| -rw-r--r-- | test/test_sliced.py | 19 |
2 files changed, 30 insertions, 1 deletions
diff --git a/test/test_backend.py b/test/test_backend.py index fd9a761..5351e52 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -296,6 +296,8 @@ def test_empty_backend(): nx.atan2(v, v) with pytest.raises(NotImplementedError): nx.transpose(M) + with pytest.raises(NotImplementedError): + nx.detach(M) def test_func_backends(nx): @@ -649,6 +651,16 @@ def test_func_backends(nx): lst_b.append(nx.to_numpy(A)) lst_name.append("transpose") + A = nx.detach(Mb) + lst_b.append(nx.to_numpy(A)) + lst_name.append("detach") + + A, B = nx.detach(Mb, Mb) + lst_b.append(nx.to_numpy(A)) + lst_name.append("detach A") + lst_b.append(nx.to_numpy(B)) + lst_name.append("detach B") + assert not nx.array_equal(Mb, vb), "array_equal (shape)" assert nx.array_equal(Mb, Mb), "array_equal (elements) - expected true" assert not nx.array_equal( diff --git a/test/test_sliced.py b/test/test_sliced.py index f54c799..7b7437a 100644 --- a/test/test_sliced.py +++ b/test/test_sliced.py @@ -10,7 +10,7 @@ import pytest import ot from ot.sliced import get_random_projections -from ot.backend import tf +from ot.backend import tf, torch def test_get_random_projections(): @@ -408,6 +408,23 @@ def test_sliced_sphere_backend_type_devices(nx): nx.assert_same_dtype_device(xb, valb) +def test_sliced_sphere_gradient(): + if torch: + import torch.nn.functional as F + + X0 = torch.randn((500, 3)) + X0 = F.normalize(X0, p=2, dim=-1) + X0.requires_grad_(True) + + X1 = torch.randn((500, 3)) + X1 = F.normalize(X1, p=2, dim=-1) + + sw = ot.sliced_wasserstein_sphere(X1, X0, n_projections=500, p=2) + grad_x0 = torch.autograd.grad(sw, X0)[0] + + assert not torch.any(torch.isnan(grad_x0)) + + def test_sliced_sphere_unif_values_on_the_sphere(): n = 100 rng = np.random.RandomState(0) |