mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
66d50060eb
Summary: Notes: 1. didn't name them as _copy_real and _copy_imag because it's desirable (but not necessary) to have these methods as tensor methods. 2. replaced old .real() and .imag() instances with _copy_real() and _copy_imag() methods 3. didn't add documentation because we plan to remove these methods when we add real and imag as tensor attributes. Pull Request resolved: https://github.com/pytorch/pytorch/pull/35879 Differential Revision: D20841760 Pulled By: anjali411 fbshipit-source-id: 7267e6fbaab9a5ce426e9396f12238994666b0dd
38 lines
1.2 KiB
Python
38 lines
1.2 KiB
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)
|
|
|
|
def test_copy_real_imag_methods(self):
|
|
real = torch.randn(4)
|
|
imag = torch.randn(4)
|
|
complex_tensor = real + 1j * imag
|
|
self.assertEqual(complex_tensor.copy_real(), real)
|
|
self.assertEqual(complex_tensor.copy_imag(), imag)
|
|
|
|
if __name__ == '__main__':
|
|
run_tests()
|