torch.utils.dlpack¶
- torch.utils.dlpack.from_dlpack(ext_tensor) Tensor [source][source]¶
Converts a tensor from an external library into a
torch.Tensor
.The returned PyTorch tensor will share the memory with the input tensor (which may have come from another library). Note that in-place operations will therefore also affect the data of the input tensor. This may lead to unexpected issues (e.g., other libraries may have read-only flags or immutable data structures), so the user should only do this if they know for sure that this is fine.
- Parameters
ext_tensor (object with
__dlpack__
attribute, or a DLPack capsule) –The tensor or DLPack capsule to convert.
If
ext_tensor
is a tensor (or ndarray) object, it must support the__dlpack__
protocol (i.e., have aext_tensor.__dlpack__
method). Otherwiseext_tensor
may be a DLPack capsule, which is an opaquePyCapsule
instance, typically produced by ato_dlpack
function or method.- Return type
Examples:
>>> import torch.utils.dlpack >>> t = torch.arange(4) # Convert a tensor directly (supported in PyTorch >= 1.10) >>> t2 = torch.from_dlpack(t) >>> t2[:2] = -1 # show that memory is shared >>> t2 tensor([-1, -1, 2, 3]) >>> t tensor([-1, -1, 2, 3]) # The old-style DLPack usage, with an intermediate capsule object >>> capsule = torch.utils.dlpack.to_dlpack(t) >>> capsule <capsule object "dltensor" at ...> >>> t3 = torch.from_dlpack(capsule) >>> t3 tensor([-1, -1, 2, 3]) >>> t3[0] = -9 # now we're sharing memory between 3 tensors >>> t3 tensor([-9, -1, 2, 3]) >>> t2 tensor([-9, -1, 2, 3]) >>> t tensor([-9, -1, 2, 3])
- torch.utils.dlpack.to_dlpack(tensor) PyCapsule ¶
Returns an opaque object (a “DLPack capsule”) representing the tensor.
Note
to_dlpack
is a legacy DLPack interface. The capsule it returns cannot be used for anything in Python other than use it as input tofrom_dlpack
. The more idiomatic use of DLPack is to callfrom_dlpack
directly on the tensor object - this works when that object has a__dlpack__
method, which PyTorch and most other libraries indeed have now.Warning
Only call
from_dlpack
once per capsule produced withto_dlpack
. Behavior when a capsule is consumed multiple times is undefined.- Parameters
tensor – a tensor to be exported
The DLPack capsule shares the tensor’s memory.