jax.dlpack.from_dlpack

Contents

jax.dlpack.from_dlpack#

jax.dlpack.from_dlpack(external_array)[source]#

Returns a Array representation of a DLPack tensor.

The returned Array shares memory with external_array.

Parameters:

external_array – an array object that has __dlpack__ and __dlpack_device__ methods, or a DLPack tensor on either CPU or GPU (legacy API).

Returns:

A jax.Array

Note

While JAX arrays are always immutable, dlpack buffers cannot be marked as immutable, and it is possible for processes external to JAX to mutate them in-place. If a jax Array is constructed from a dlpack buffer and the buffer is later modified in-place, it may lead to undefined behavior when using the associated JAX array.