mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
b004307252
Test Plan: manual inspection & sandcastle Reviewed By: zertosh Differential Revision: D30279364 fbshipit-source-id: c1ed77dfe43a3bde358f92737cd5535ae5d13c9a
37 lines
1.2 KiB
Python
37 lines
1.2 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.0, 3.0 + 5.0j], 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()
|