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
|
# Owner(s): ["module: distributions"]
import pytest
import torch
from torch.distributions.utils import tril_matrix_to_vec, vec_to_tril_matrix
@pytest.mark.parametrize('shape', [
(2, 2),
(3, 3),
(2, 4, 4),
(2, 2, 4, 4),
])
def test_tril_matrix_to_vec(shape):
mat = torch.randn(shape)
n = mat.shape[-1]
for diag in range(-n, n):
actual = mat.tril(diag)
vec = tril_matrix_to_vec(actual, diag)
tril_mat = vec_to_tril_matrix(vec, diag)
assert torch.allclose(tril_mat, actual)
if __name__ == '__main__':
pytest.main([__file__])
|
