pearu / pydlpack

DLPack data exchange interface in Python

pydlpack

PyDLPack is a Python library for exchanging data between different array libraries using DLPack: Open In Memory Tensor Structure. The provider library does not need to implement the DLPack support, it will be sufficent if the provider library implements one of the following protocols:

• dmlc/dlpack
• Array Interface Protocol, version 3
• CUDA Array Interface, version 3
• Python Buffer protocol

Array Interface Protocol, version 2 support can be provided on request.

Currently, the package is tested with the following consumers

• NumPy, numpy.from_dlpack
• PyTorch, torch.from_dlpack
• CuPy, cupy.from_dlpack
• Jax, jax.numpy.from_dlpack
• Tensorflow, tf.experimental.dlpack.from_dlpack
• cuDF, cudf.from_dlpack

using the following provider objects with devices:

• Numpy ndarray, CPU
• Torch Tensor, CPU and CUDA
• CuPy ndarray, CUDA
• Numba DeviceNDArray, CUDA
• Jax Array, CPU and CUDA
• Tensorflow Tensor, CPU
• Python bytes, CPU
• Python bytearray, CPU
• Python array, CPU
• Numpy memmap, CPU
• Python mmap, CPU

Install

• pydlpack in conda-forge

  conda install pydlpack

• pydlpack in PyPi:

  pip install pydlpack

Basic usage

>>> from dlpack import asdlpack
>>> import torch
>>> dl = asdlpack(b"Hello!")
>>> torch.from_dlpack(dl)
tensor([ 72, 101, 108, 108, 111,  33], dtype=torch.uint8)

that is, the Python package dlpack provides a function asdlpack that input can be any object that implements one of the above mentioned protocols and it will return a light-weight DLPackObject instance which implements the DLPack protocol methods __dlpack__(stream=None) and __dlpack_device__(). This DLPackObject instance can be used as an argument to a consumer.from_dlpack(obj) function of any DLPack-compatible consumer library (a partial list of such libraries is listed above). For example:

>>> from dlpack import asdlpack
>>> import numba.cuda
>>> import numpy
>>> arr = numba.cuda.to_device(numpy.array([[1, 2], [3, 4]]))
>>> arr
<numba.cuda.cudadrv.devicearray.DeviceNDArray object at 0x7fbed9c548b0>
>>> dl = asdlpack(arr)
>>> import torch
>>> torch.from_dlpack(dl)
tensor([[1, 2],
        [3, 4]], device='cuda:0')
>>> import jax
>>> jax.numpy.from_dlpack(dl)
Array([[1, 2],
       [3, 4]], dtype=int32)
>>> import cupy
>>> cupy.from_dlpack(dl)
array([[1, 2],
       [3, 4]])

that is, the DLPackObject instance can be efficiently used for exchanging the CUDA buffer created using Numba to_device functionality with different consumer objects such as torch.Tensor, jax.Array, and cupy.ndarray while all these array objects share the same CUDA memory.

Testing

It is a non-trivial task to install all dlpack-compatible libraries into the same environment. Therefore, dlpack tests are included the dlpack package so that one can import dlpack and run the tests on DLPack-compatible objects that are available in a particular environment. For example:

>>> import dlpack.tests
>>> dlpack.tests.run()