Test reading __cuda_array_interface__ inferred strides. (pytorch#31451)

Summary: Pull Request resolved: pytorch#31451 The PR that fixed this, pytorch#24947, didn't add a test. Fixes: pytorch#31443 Test Plan: Imported from OSS Differential Revision: D19170020 Pulled By: gchanan fbshipit-source-id: bdbf09989ac8a61b1b70bb1ddee103caa8ef435b

Author: gchanan

test/test_numba_integration.py

@@ -309,6 +309,30 @@ def test_from_cuda_array_interface(self):
  torch_ary += 42
  self.assertEqual(torch_ary.cpu().data.numpy(), numpy.asarray(numba_ary) + 42)
 
+ @unittest.skipIf(not TEST_NUMPY, "No numpy")
+ @unittest.skipIf(not TEST_CUDA, "No cuda")
+ @unittest.skipIf(not TEST_NUMBA_CUDA, "No numba.cuda")
+ def test_from_cuda_array_interface_inferred_strides(self):
+ """torch.as_tensor(numba_ary) should have correct inferred (contiguous) strides"""
+ # This could, in theory, be combined with test_from_cuda_array_interface but that test
+ # is overly strict: it checks that the exported protocols are exactly the same, which
+ # cannot handle differing exported protocol versions.
+ dtypes = [
+ numpy.float64,
+ numpy.float32,
+ numpy.int64,
+ numpy.int32,
+ numpy.int16,
+ numpy.int8,
+ numpy.uint8,
+ ]
+ for dtype in dtypes:
+ numpy_ary = numpy.arange(6).reshape(2, 3).astype(dtype),
+ numba_ary = numba.cuda.to_device(numpy_ary)
+ self.assertTrue(numba_ary.is_c_contiguous())
+ torch_ary = torch.as_tensor(numba_ary, device="cuda")
+ self.assertTrue(torch_ary.is_contiguous())
+
  @unittest.skipIf(not TEST_NUMPY, "No numpy")
  @unittest.skipIf(not TEST_CUDA, "No cuda")
  @unittest.skipIf(not TEST_NUMBA_CUDA, "No numba.cuda")