DLPack

Via a DLPack capsule, you can borrow a tensor from external software as a nnabla.NdArray and can share a NdArray to external software.

nnabla.utils.dlpack.from_dlpack(dlp, arr=None)

Decode a DLPack to NdArray.

Example:

# Create a tensor of an external tool, and encode as an DLPack.
import torch
from torch.utils.dlpack import to_dlpack
t = torch.ones((5, 5), dtype=torch.float32,
               device=torch.device('cuda'))
dlp = to_dlpack(t)

# Borrow the DLPack tensor as nnabla.NdArray
from nnabla.utils.dlpack import from_dlpack
arr = from_dlpack(dlp)

If you want to move an ownership of DLPack to an exiting NdArray;

from nnabla import NdArray
arr = NdArray()
from_dlpack(dlp, arr=arr)
Parameters:
  • dlp (PyCapsule) – A PyCapsule object of a DLManagedTensor (as "dltensor") which internal memory is borrowed by a tensor of an external package. The ownership of the DLManagedTensor is moved to an NdArray object, and the PyCapsule object is marked as "used_dltensor" to inform that the ownership has been moved.

  • arr (NdArray) – If specified, a given DLPack is decoded to it. Otherwise, it creates a new NdArray object and decodes the DLPack to it.

Returns:

an NdArray object borrowing the DLPack tensor.

Return type:

NdArray

nnabla.utils.dlpack.to_dlpack(a, dtype=None, ctx=None)

Returns a DLPack which owns an internal array object borrowed by a specified NdArray.

Example:

# Create a nnabla.NdArray in CUDA.
import nnabla as nn
from nnabla.ext_utils import get_extension_context
ctx = get_extension_context('cudnn')
nn.set_default_context(ctx)

a = nn.NdArray.from_numpy_array(np.ones((5, 5), dtype=np.float32))
a.cast(np.float32, ctx)

# Expose as a DLPack.
from nnabla.utils.dlpack import to_dlpack
dlp = to_dlpack(a)

# Use the DLPack in PyTorch.
import torch
from torch.utils.dlpack import from_dlpack
t = from_dlpack(dlp)

# Changing the values in Torch will also be affected in nnabla
# because they share memory.
t.add_(1)
print(a.data)  # All values become 2.
Parameters:
  • a (NdArray) – An NdArray object. An internal array which is recently modified or created will be encoded into a DLPack.

  • dtype (numpy.dtype) – If specified, in-place cast operation may be performed before encoding it to a DLPack.

  • ctx (Context) – If specified, in-place device transfer operation may be performed before encoding it into a DLPack.

Returns:

A PyCapsule object of a DLManagedTensor (as "dltensor") which internal memory is borrowed by the specified NdArray.

Return type:

PyCapsule