mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
266e278ccf
Partiall addresses #123062 UFMT formatting on - test/distributions - test/error_messages, test/forward_backward_compatability Pull Request resolved: https://github.com/pytorch/pytorch/pull/123527 Approved by: https://github.com/huydhn
31 lines
670 B
Python
31 lines
670 B
Python
# Owner(s): ["module: distributions"]
|
|
|
|
import pytest
|
|
|
|
import torch
|
|
from torch.distributions.utils import tril_matrix_to_vec, vec_to_tril_matrix
|
|
from torch.testing._internal.common_utils import run_tests
|
|
|
|
|
|
@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__":
|
|
run_tests()
|