Updated dlpack_test

Author: oandreeva-nv

qa/python_models/dlpack_test/model.py

@@ -1,4 +1,4 @@
-# Copyright 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2021-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -56,6 +56,23 @@ def test_pytorch_dlpack(self):
             self.assertTrue(
                 pytorch_tensor.type() == pytorch_tensor_dlpack.type())
 
+            # Now let's check that upgraded DLPack implementation also
+            # works as expected, i.e. from_dlpack should work with 
+            # external pytorch tensor directly
+
+            pb_tensor_upgraded = pb_utils.Tensor.from_dlpack('test_tensor',
+                                                    pytorch_tensor)
+            self.assertTrue(
+                np.all(pb_tensor_upgraded.as_numpy() == pytorch_tensor.numpy()))
+
+            # Here we check that `pb_tensor` as a producer, properly
+            # invokes `__dlpack__` and `__dlpack_device__`
+            pytorch_tensor_dlpack = from_dlpack(pb_tensor_upgraded)
+            self.assertTrue(torch.all(pytorch_tensor_dlpack == pytorch_tensor))
+
+            self.assertTrue(
+                pytorch_tensor.type() == pytorch_tensor_dlpack.type())  
+
     def test_non_contiguous_error(self):
         pytorch_tensor = torch.rand([20, 30], dtype=torch.float16)
 
@@ -83,6 +100,8 @@ def test_dlpack_string_tensor(self):
 
     def test_dlpack_gpu_tensors(self):
         # Test different dtypes
+        # PyTorch does not support DLPack bool type yet:
+        # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/DLConvertor.cpp
         pytorch_dtypes = [
             torch.float16, torch.float32, torch.float64, torch.int8,
             torch.int16, torch.int32, torch.int64, torch.uint8
@@ -100,27 +119,50 @@ def test_dlpack_gpu_tensors(self):
             # the same
             pytorch_tensor_dlpack = from_dlpack(pb_tensor.to_dlpack())
             self.assertTrue(torch.all(pytorch_tensor_dlpack == pytorch_tensor))
-
-            # DLPack does not properly support bool type:
-            # https://github.com/google/jax/issues/4719
-            if pytorch_dtype != torch.bool:
-                self.assertTrue(
-                    pytorch_tensor.type() == pytorch_tensor_dlpack.type())
-            else:
-                self.assertFalse(
-                    pytorch_tensor.type() == pytorch_tensor_dlpack.type())
+            self.assertTrue(
+                pytorch_tensor.type() == pytorch_tensor_dlpack.type())
+            
+            # Now we make sure that updated DLPack implementation works
+            # with GPU as well
+            pb_tensor = pb_utils.Tensor.from_dlpack('test_tensor',
+                                                    pytorch_tensor)
+            pytorch_tensor_dlpack = from_dlpack(pb_tensor)
+            self.assertTrue(torch.all(pytorch_tensor_dlpack == pytorch_tensor))
+            self.assertTrue(
+                pytorch_tensor.type() == pytorch_tensor_dlpack.type())
+            
 
     def test_dlpack_gpu_numpy(self):
         # DLPack tesnors that are in GPU cannot be converted to NumPy
         pytorch_tensor = torch.rand([100], dtype=torch.float16,
                                     device='cuda') * 100
         pb_tensor = pb_utils.Tensor.from_dlpack('tensor',
                                                 to_dlpack(pytorch_tensor))
+        # Make sure that `__dlpack_device__` works as expected
+        self.assertTrue(pb_tensor.__dlpack_device__() == pytorch_tensor.__dlpack_device__())
+    
         with self.assertRaises(Exception) as e:
             pb_tensor.as_numpy()
         self.assertTrue(
             str(e.exception) ==
             'Tensor is stored in GPU and cannot be converted to NumPy.')
+    
+    def test_dlpack_cpu_numpy(self):
+        # Check compatibiity of PbTensor DLPack implementation
+        # with numpy
+        pytorch_tensor = torch.rand([100], dtype=torch.float16,
+                                    device='cpu') * 100
+        pb_tensor = pb_utils.Tensor.from_dlpack('tensor', pytorch_tensor)
+        numpy_tensor_dlpack = np.from_dlpack(pb_tensor)
+        self.assertTrue(np.all(numpy_tensor_dlpack == pytorch_tensor.numpy()))
+        # Make sure that `__dlpack_device__` works as expected
+        self.assertTrue(pb_tensor.__dlpack_device__() == pytorch_tensor.__dlpack_device__())
+    
+    def test_pdtensor_bool_internal_support(self):
+        bool_array = np.asarray([False, True])
+        bool_tensor = pb_utils.Tensor('tensor', bool_array)
+        bool_tensor_dlpack = pb_utils.Tensor.from_dlpack('tensor', bool_tensor)
+        self.assertTrue(np.all(bool_array == bool_tensor_dlpack.as_numpy()))
 
 
 class TritonPythonModel: