mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
b3cdec88e3
Summary: There was a permutation operation missing in each of the complex vector files. I also added some test cases, the last two of which fail under the current implementation. This PR fixes that: all the testcases pass. Fixes https://github.com/pytorch/pytorch/issues/35532 dylanbespalko Pull Request resolved: https://github.com/pytorch/pytorch/pull/35715 Differential Revision: D20857024 Pulled By: anjali411 fbshipit-source-id: 4eecd8f0863faa838300951626f26b89e6cc9c6b
32 lines
963 B
Python
32 lines
963 B
Python
import math
|
|
import torch
|
|
from torch.testing._internal.common_utils import TestCase, run_tests, TEST_NUMPY
|
|
import unittest
|
|
|
|
if TEST_NUMPY:
|
|
import numpy as np
|
|
|
|
devices = (torch.device('cpu'), torch.device('cuda:0'))
|
|
|
|
|
|
class TestComplexTensor(TestCase):
|
|
def test_to_list_with_complex_64(self):
|
|
# test that the complex float tensor has expected values and
|
|
# there's no garbage value in the resultant list
|
|
self.assertEqual(torch.zeros((2, 2), dtype=torch.complex64).tolist(), [[0j, 0j], [0j, 0j]])
|
|
|
|
@unittest.skipIf(not TEST_NUMPY, "Numpy not found")
|
|
def test_exp(self):
|
|
def exp_fn(dtype):
|
|
a = torch.tensor(1j, dtype=dtype) * torch.arange(18) / 3 * math.pi
|
|
expected = np.exp(a.numpy())
|
|
actual = torch.exp(a)
|
|
self.assertEqual(actual, torch.from_numpy(expected))
|
|
|
|
exp_fn(torch.complex64)
|
|
exp_fn(torch.complex128)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
run_tests()
|