forked from probabilists/zuko
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest_transforms.py
120 lines (84 loc) · 3.38 KB
/
test_transforms.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
r"""Tests for the zuko.transforms module."""
import pytest
import torch
from torch import randn
from torch.distributions import *
from zuko.transforms import *
@pytest.mark.parametrize("batched", [False, True])
def test_univariate_transforms(batched: bool):
batch = (256,) if batched else ()
ts = [
IdentityTransform(),
CosTransform(),
SinTransform(),
SoftclipTransform(),
CircularShiftTransform(),
SignedPowerTransform(randn(batch)),
MonotonicAffineTransform(randn(batch), randn(batch)),
MonotonicRQSTransform(randn(*batch, 8), randn(*batch, 8), randn(*batch, 7)),
MonotonicTransform(lambda x: x**3),
BernsteinTransform(randn(*batch, 16)),
BoundedBernsteinTransform(randn(*batch, 16)),
GaussianizationTransform(randn(*batch, 8), randn(*batch, 8)),
UnconstrainedMonotonicTransform(lambda x: torch.exp(-(x**2)) + 1e-2, randn(batch)),
SOSPolynomialTransform(randn(*batch, 3, 5), randn(batch)),
]
for t in ts:
# Call
if hasattr(t.domain, "lower_bound"):
x = torch.linspace(t.domain.lower_bound + 1e-2, t.domain.upper_bound - 1e-2, 256)
else:
x = torch.linspace(-5.0, 5.0, 256)
y = t(x)
assert x.shape == y.shape, t
# Inverse
z = t.inv(y)
assert torch.allclose(x, z, atol=1e-4), t
# Jacobian
J = torch.autograd.functional.jacobian(t, x)
assert (torch.triu(J, diagonal=+1) == 0).all(), t
assert (torch.tril(J, diagonal=-1) == 0).all(), t
ladj = torch.diag(J).abs().log()
assert torch.allclose(t.log_abs_det_jacobian(x, y), ladj, atol=1e-4), t
# Compound call
y_comp, ladj_comp = t.call_and_ladj(x)
assert torch.allclose(y_comp, y, atol=1e-4), t
assert torch.allclose(ladj_comp, ladj, atol=1e-4), t
# Inverse Jacobian
J = torch.autograd.functional.jacobian(t.inv, y)
assert (torch.triu(J, diagonal=+1) == 0).all(), t
assert (torch.tril(J, diagonal=-1) == 0).all(), t
ladj = torch.diag(J).abs().log()
assert torch.allclose(t.inv.log_abs_det_jacobian(y, z), ladj, atol=1e-4), t
def test_multivariate_transforms():
A, B = torch.randn(5, 16), torch.randn(16, 5)
f = lambda t, x: torch.sigmoid(x @ A) @ B
ts = [
FreeFormJacobianTransform(f, 0.0, 1.0, atol=1e-7, rtol=1e-6),
PermutationTransform(torch.randperm(5)),
RotationTransform(randn(5, 5)),
LULinearTransform(randn(5, 5)),
]
for t in ts:
# Shapes
x = randn(256, 5)
y = t(x)
assert x.shape == y.shape, t
# Inverse
z = t.inv(y)
assert torch.allclose(x, z, atol=1e-4), t
# Jacobian
x = randn(5)
y = t(x)
J = torch.autograd.functional.jacobian(t, x)
ladj = torch.linalg.slogdet(J).logabsdet
assert torch.allclose(t.log_abs_det_jacobian(x, y), ladj, atol=1e-4), t
# Compound call
y_comp, ladj_comp = t.call_and_ladj(x)
assert torch.allclose(y_comp, y, atol=1e-4), t
assert torch.allclose(ladj_comp, ladj, atol=1e-4), t
# Inverse Jacobian
z = t.inv(y)
J = torch.autograd.functional.jacobian(t.inv, y)
ladj = torch.linalg.slogdet(J).logabsdet
assert torch.allclose(t.inv.log_abs_det_jacobian(y, z), ladj, atol=1e-4), t