Files
pytorch/test/test_complex.py
anjali411 f9f22c8b5c Add serialization logic for complex numbers (#51287)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51287

This reverts commit dfdb1547b9c1934904bfd137b4007d6a46a6f597.

Test Plan: Imported from OSS

Reviewed By: SplitInfinity

Differential Revision: D26131165

Pulled By: anjali411

fbshipit-source-id: 047167fac594ddb670c5e169446e90e74991679a
2021-01-28 17:25:35 -08:00

27 lines
1.1 KiB
Python

import torch
from torch.testing._internal.common_device_type import instantiate_device_type_tests, dtypes
from torch.testing._internal.common_utils import TestCase, run_tests
devices = (torch.device('cpu'), torch.device('cuda:0'))
class TestComplexTensor(TestCase):
@dtypes(*torch.testing.get_all_complex_dtypes())
def test_to_list(self, device, dtype):
# 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), device=device, dtype=dtype).tolist(), [[0j, 0j], [0j, 0j]])
@dtypes(torch.float32, torch.float64)
def test_dtype_inference(self, device, dtype):
# issue: https://github.com/pytorch/pytorch/issues/36834
default_dtype = torch.get_default_dtype()
torch.set_default_dtype(dtype)
x = torch.tensor([3., 3. + 5.j], device=device)
torch.set_default_dtype(default_dtype)
self.assertEqual(x.dtype, torch.cdouble if dtype == torch.float64 else torch.cfloat)
instantiate_device_type_tests(TestComplexTensor, globals())
if __name__ == '__main__':
run_tests()