DI-engine/ding/torch_utils/tests/test_parameter.py

import unittest
import pytest
import torch
from ding.torch_utils.parameter import NonegativeParameter, TanhParameter


@pytest.mark.unittest
def test_nonegative_parameter():
    nonegative_parameter = NonegativeParameter(torch.tensor([2.0, 3.0]))
    assert torch.sum(torch.abs(nonegative_parameter() - torch.tensor([2.0, 3.0]))) == 0
    nonegative_parameter.set_data(torch.tensor(1))
    assert nonegative_parameter() == 1


@pytest.mark.unittest
def test_tanh_parameter():
    tanh_parameter = TanhParameter(torch.tensor([0.5, -0.2]))
    assert torch.isclose(tanh_parameter() - torch.tensor([0.5, -0.2]), torch.zeros(2), atol=1e-6).all()
    tanh_parameter.set_data(torch.tensor(0.3))
    assert tanh_parameter() == 0.3


if __name__ == "__main__":
    test_nonegative_parameter()
    test_tanh_parameter()