A wrapper class that holds references to the PyCapsule and its associated data. This class prevents Python's garbage collector from collecting the shape_array and managed_tensor objects while the capsule is still in use. It serves as a container to maintain the lifecycle of all DLP
| 84 | |
| 85 | # Wrapper class to prevent Python garbage collection of DLPack-related objects |
| 86 | class CapsuleWrapper: |
| 87 | """ |
| 88 | A wrapper class that holds references to the PyCapsule and its associated data. |
| 89 | |
| 90 | This class prevents Python's garbage collector from collecting the shape_array and |
| 91 | managed_tensor objects while the capsule is still in use. It serves as a container |
| 92 | to maintain the lifecycle of all DLPack-related objects. |
| 93 | """ |
| 94 | |
| 95 | def __init__(self, capsule, shape_array, managed_tensor): |
| 96 | """ |
| 97 | Initialize the CapsuleWrapper with the necessary objects. |
| 98 | |
| 99 | Parameters: |
| 100 | capsule: The PyCapsule object that follows the DLPack protocol |
| 101 | shape_array: The array containing tensor shape information |
| 102 | managed_tensor: The DLManagedTensor instance that the capsule points to |
| 103 | """ |
| 104 | self.capsule = capsule # The main PyCapsule object that can be passed to other libraries |
| 105 | self._shape_array = shape_array # Keep reference to prevent garbage collection |
| 106 | self._managed_tensor = managed_tensor # Keep reference to prevent garbage collection |
| 107 | |
| 108 | |
| 109 | def create_dlpack_capsule(ptr, segment_size, segment_stride, num_segments, torch_dtype, dev_id): |
no outgoing calls
no test coverage detected