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

Co-authored-by: iefgnoix <isaacwxf23@gmail.com>

Author: ysiraichi

.torch_pin

@@ -0,0 +1 @@
+#138470

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):