mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
57d01be92b
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/38102 Test Plan: Imported from OSS Differential Revision: D21477060 Pulled By: VitalyFedyunin fbshipit-source-id: 25e0fd837ca9bfccf0ce994c80f7790c894096d4
46 lines
1.6 KiB
Python
46 lines
1.6 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
|
|
# TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
|
|
self.assertEqualIgnoreType(complex_tensor.copy_real(), real)
|
|
# TODO(#38095): Replace assertEqualIgnoreType. See issue #38095
|
|
self.assertEqualIgnoreType(complex_tensor.copy_imag(), imag)
|
|
|
|
def test_dtype_inference(self):
|
|
# issue: https://github.com/pytorch/pytorch/issues/36834
|
|
torch.set_default_dtype(torch.double)
|
|
x = torch.tensor([3., 3. + 5.j])
|
|
self.assertEqual(x.dtype, torch.cdouble)
|
|
|
|
if __name__ == '__main__':
|
|
run_tests()
|