Update tests

pytorch · Jul 15, 2023 · 19bc63a · 19bc63a
1 parent 0e3d5b6
commit 19bc63a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/test/test_prims.py b/test/test_prims.py
@@ -1212,6 +1212,16 @@ def test_unbind(self):
         expect = torch.unbind(a, 1)
         self.assertEqual(actual, expect)
 
+    def test_logspace_with_complex_input(self):
+        actual = refs.logspace(2, 10 + 5j, steps=5)
+        expect = torch.logspace(2, 10 + 5j, steps=5)
+        self.assertEqual(actual, expect)
+
+    def test_linspace_with_complex_input(self):
+        actual = refs.linspace(2, 10 + 5j, steps=5)
+        expect = torch.linspace(2, 10 + 5j, steps=5)
+        self.assertEqual(actual, expect)
+
 
 instantiate_device_type_tests(TestRefs, globals())
 

diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
@@ -1002,7 +1002,6 @@ def sample_inputs_linspace(op, device, dtype, requires_grad, **kwargs):
         yield SampleInput(start, args=(end, nstep), kwargs={"dtype": dtype, "device": device})
 
     yield SampleInput(1, args=(3, 1))
-    yield SampleInput(1 + 2j, args=(10 + 5j, 5))  # inputs of complex type
 
 
 def sample_inputs_logspace(op, device, dtype, requires_grad, **kwargs):
@@ -1022,7 +1021,6 @@ def sample_inputs_logspace(op, device, dtype, requires_grad, **kwargs):
             yield SampleInput(start, args=(end, nstep, base), kwargs={"dtype": dtype, "device": device})
 
     yield SampleInput(1, args=(3, 1, 2.))
-    yield SampleInput(1 + 2j, args=(10 + 5j, 5))  # inputs of complex type
 
 
 def sample_inputs_isclose(op, device, dtype, requires_grad, **kwargs):