jax.dlpack module#

from_dlpack(dlpack)

Returns a DeviceArray representation of a DLPack tensor.

to_dlpack(x[, take_ownership])

Returns a DLPack tensor that encapsulates a DeviceArray x.