Spaces:
Sleeping
Sleeping
| import pytest | |
| import torch | |
| from ding.torch_utils.distribution import Pd, CategoricalPd, CategoricalPdPytorch | |
| class TestProbDistribution: | |
| def test_Pd(self): | |
| pd = Pd() | |
| with pytest.raises(NotImplementedError): | |
| pd.neglogp(torch.randn(5, )) | |
| with pytest.raises(NotImplementedError): | |
| pd.noise_mode() | |
| with pytest.raises(NotImplementedError): | |
| pd.mode() | |
| with pytest.raises(NotImplementedError): | |
| pd.sample() | |
| def test_CatePD(self): | |
| pd = CategoricalPd() | |
| logit1 = torch.randn(3, 5, requires_grad=True) | |
| logit2 = torch.randint(5, (3, ), dtype=torch.int64) | |
| pd.update_logits(logit1) | |
| entropy = pd.neglogp(logit2) | |
| assert entropy.requires_grad | |
| assert entropy.shape == torch.Size([]) | |
| entropy = pd.entropy() | |
| assert entropy.requires_grad | |
| assert entropy.shape == torch.Size([]) | |
| entropy = pd.entropy(reduction=None) | |
| assert entropy.requires_grad | |
| assert entropy.shape == torch.Size([3]) | |
| ret = pd.sample() | |
| assert ret.shape == torch.Size([3]) | |
| ret = pd.sample(viz=True) | |
| assert ret[0].shape == torch.Size([3]) | |
| ret = pd.mode() | |
| assert ret.shape == torch.Size([3]) | |
| ret = pd.mode(viz=True) | |
| assert ret[0].shape == torch.Size([3]) | |
| ret = pd.noise_mode() | |
| assert ret.shape == torch.Size([3]) | |
| ret = pd.noise_mode(viz=True) | |
| assert ret[0].shape == torch.Size([3]) | |
| pd = CategoricalPdPytorch() | |
| pd.update_logits(logit1) | |
| ret = pd.sample() | |
| assert ret.shape == torch.Size([3]) | |
| ret = pd.mode() | |
| assert ret.shape == torch.Size([3]) | |
| entropy = pd.entropy(reduction='mean') | |
| assert entropy.requires_grad | |
| assert entropy.shape == torch.Size([]) | |
| entropy = pd.entropy(reduction=None) | |
| assert entropy.requires_grad | |
| assert entropy.shape == torch.Size([3]) | |