Add complex tensor dtypes for the __cuda_array_interface__ spec (pytorch#42918)

Summary:
Fixes pytorch#42860

The `__cuda_array_interface__` tensor specification is missing the appropriate datatypes for the newly merged complex64 and complex128 tensors. This PR addresses this issue by casting:

* `torch.complex64` to 'c8'
* `torch.complex128` to 'c16'

Pull Request resolved: pytorch#42918

Reviewed By: izdeby

Differential Revision: D23130219

Pulled By: anjali411

fbshipit-source-id: 5f8ee8446a71cad2f28811afdeae3a263a31ad11

Author: awthomp

test/test_numba_integration.py

@@ -105,6 +105,8 @@ def test_array_adaptor(self):
         """Torch __cuda_array_adaptor__ exposes tensor data to numba.cuda."""
 
         torch_dtypes = [
+            torch.complex64,
+            torch.complex128,
             torch.float16,
             torch.float32,
             torch.float64,
@@ -244,6 +246,8 @@ def test_from_cuda_array_interface(self):
         """
 
         dtypes = [
+            numpy.complex64,
+            numpy.complex128,
             numpy.float64,
             numpy.float32,
             numpy.int64,
@@ -263,31 +267,31 @@ def test_from_cuda_array_interface(self):
                 numba_ary = numba.cuda.to_device(numpy_ary)
                 torch_ary = torch.as_tensor(numba_ary, device="cuda")
                 self.assertEqual(numba_ary.__cuda_array_interface__, torch_ary.__cuda_array_interface__)
-                self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary))
+                self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary, dtype=dtype))
 
                 # Check that `torch_ary` and `numba_ary` points to the same device memory
                 torch_ary += 42
-                self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary))
+                self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary, dtype=dtype))
 
             # Implicit-copy because `torch_ary` is a CPU array
             for numpy_ary in numpy_arys:
                 numba_ary = numba.cuda.to_device(numpy_ary)
                 torch_ary = torch.as_tensor(numba_ary, device="cpu")
-                self.assertEqual(torch_ary.data.numpy(), numpy.asarray(numba_ary))
+                self.assertEqual(torch_ary.data.numpy(), numpy.asarray(numba_ary, dtype=dtype))
 
                 # Check that `torch_ary` and `numba_ary` points to different memory
                 torch_ary += 42
-                self.assertEqual(torch_ary.data.numpy(), numpy.asarray(numba_ary) + 42)
+                self.assertEqual(torch_ary.data.numpy(), numpy.asarray(numba_ary, dtype=dtype) + 42)
 
             # Explicit-copy when using `torch.tensor()`
             for numpy_ary in numpy_arys:
                 numba_ary = numba.cuda.to_device(numpy_ary)
                 torch_ary = torch.tensor(numba_ary, device="cuda")
-                self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary))
+                self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary, dtype=dtype))
 
                 # Check that `torch_ary` and `numba_ary` points to different memory
                 torch_ary += 42
-                self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary) + 42)
+                self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary, dtype=dtype) + 42)
 
     @unittest.skipIf(not TEST_NUMPY, "No numpy")
     @unittest.skipIf(not TEST_CUDA, "No cuda")

torch/tensor.py

@@ -690,6 +690,8 @@ def __cuda_array_interface__(self):
         # CUDA devices are little-endian and tensors are stored in native byte
         # order. 1-byte entries are endian-agnostic.
         typestr = {
+            torch.complex64: "<c8",
+            torch.complex128: "<c16",
             torch.float16: "<f2",
             torch.float32: "<f4",
             torch.float64: "<f8",