Skip to main content

Module dlpack

Module dlpack 

Source
Expand description

DLPack tensor interop for scirs2-python

Provides from_dlpack and to_dlpack entry points that follow the DLPack 1.0 protocol. Full zero-copy sharing with PyTorch, JAX, CuPy, TensorFlow etc. requires the calling Python environment to have the relevant library installed; the Rust side handles the capsule protocol.

§DLPack protocol

A DLPack capsule is a PyCapsule object whose name is "dltensor". After the consumer takes ownership, the capsule is renamed to "used_dltensor" so double-frees are prevented.

§Python usage

import torch
import scirs2

t = torch.randn(3, 4)
# PyTorch tensors expose __dlpack__() / __dlpack_device__()
capsule = t.__dlpack__()
arr = scirs2.from_dlpack(capsule)   # -> scirs2 array (NumPy-compatible)

# Round-trip: export back
cap2 = scirs2.to_dlpack(arr)
t2 = torch.from_dlpack(cap2)

Functions§

from_dlpack
Convert a DLPack capsule (from PyTorch, JAX, CuPy, TensorFlow, …) into a scirs2 NumPy-compatible array.
register_dlpack_module
Register DLPack interop functions on the given module.
to_dlpack
Export a scirs2 (NumPy-compatible) array as a DLPack PyCapsule.