DLPack: add test using PyTorch DLPack functions. (#8294)

Co-authored-by: iefgnoix <isaacwxf23@gmail.com>
pytorch · Oct 30, 2024 · dc20b2d · dc20b2d
1 parent 8177447
commit dc20b2d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/.torch_pin b/.torch_pin
@@ -0,0 +1 @@
+#138470
diff --git a/test/test_operations.py b/test/test_operations.py
@@ -2912,6 +2912,17 @@ def test_dlpack_xla_to_pytorch_cuda(self):
     cuda_t1[0] = cuda_t1[0] + 20
     self.assertTrue(torch.allclose(xla_t1.cpu(), cuda_t1.cpu()))
 
+  @onlyIfTorchSupportsCUDA
+  @onlyIfPJRTDeviceIsCUDA
+  def test_dlpack_xla_to_pytorch_cuda_protocol_conversion(self):
+    xla_t1 = torch.arange(5).to(xm.xla_device())
+    caps_t1 = torch.utils.dlpack.to_dlpack(xla_t1)
+    cuda_t1 = torch.utils.dlpack.from_dlpack(caps_t1)
+    self.assertEqual(cuda_t1.device.type, 'cuda')
+    self.assertEqual(cuda_t1.device.index, xla_t1.device.index)
+    cuda_t1[0] = cuda_t1[0] + 20
+    self.assertTrue(torch.allclose(xla_t1.cpu(), cuda_t1.cpu()))
+
   @onlyIfTorchSupportsCUDA
   @onlyIfPJRTDeviceIsCUDA
   def test_dlpack_non_default_layout(self):