DLPack data exchange interface in Python
Project description
pydlpack
PyDLPack is a Python library for exchanging data between different array libraries using [DLPack: Open In Memory Tensor Structure]((https://github.com/dmlc/dlpack/). 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
-
conda install pydlpack
-
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()
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
File details
Details for the file pydlpack-0.2.1.tar.gz
.
File metadata
- Download URL: pydlpack-0.2.1.tar.gz
- Upload date:
- Size: 22.2 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.9.18
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | 0b3b1b0dd2292fac72a020b8c191fdaf32be6f9471ec696c88f753fbc042c28f |
|
MD5 | 17196504d7f25cb2c6b23ac319363766 |
|
BLAKE2b-256 | 872348c707f166468bf116740b2922b5ff745fdfd72bb40304a0358ce495d4b4 |
File details
Details for the file pydlpack-0.2.1-py3-none-any.whl
.
File metadata
- Download URL: pydlpack-0.2.1-py3-none-any.whl
- Upload date:
- Size: 19.4 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/4.0.2 CPython/3.9.18
File hashes
Algorithm | Hash digest | |
---|---|---|
SHA256 | b42ef189511e5ab4e1eafa42c5f0825d2b4349518020a3a9c33390890deb7cb2 |
|
MD5 | d4fadb5f5cd11f9cde4a320bfa73e013 |
|
BLAKE2b-256 | 3df8ae765c0f3fac5f10119d51c9b82002ea709c1899190c4743e0b124e82149 |