torch-c-dlpack-ext

import torch # WARNING: This library is primarily an internal utility of PyTorch # and is not intended for direct use by external users. The API may change # without warning and may not follow semantic versioning for external consumption. try: from torch_c_dlpack_ext import to_dlpack, from_dlpack except ImportError: print("ERROR: torch-c-dlpack-ext not installed or failed to import.") print("This library is an internal utility and not meant for direct use.") print("Please ensure 'torch' is also installed.") exit(1) print("Successfully imported torch_c_dlpack_ext (internal utility).") # 1. Create a PyTorch tensor tensor_orig = torch.randn(2, 3) print(f"\nOriginal PyTorch Tensor:\n{tensor_orig}") # 2. Convert to DLPack capsule dlpack_capsule = to_dlpack(tensor_orig) print(f"\nDLPack Capsule (type): {type(dlpack_capsule)}") # 3. Convert back to a PyTorch tensor tensor_restored = from_dlpack(dlpack_capsule) print(f"\nRestored PyTorch Tensor (from DLPack):\n{tensor_restored}") # 4. Verify data (DLPack shares underlying memory, so data should be identical) print(f"\nAre restored and original tensor data close? {torch.allclose(tensor_orig, tensor_restored)}") print(f"Are restored and original tensors the exact same Python object? {tensor_orig is tensor_restored}") # Expected output for last line: False, as 'from_dlpack' creates a new Python Tensor object # sharing the same underlying data.