-
Hi, I wanted to use jax together with moderngl (an opengl wrapper for python) to build an app that allows the user to visualize and interact with images, but also leverages jax to do various image processing. moderngl offers the Buffer as a handle for GPU memory, but as far as I know it's not possible to instantiate a I can follow the route of:
But if possible I'd like to avoid the GPU mem -> system mem -> GPU mem roundtrip. Thanks very much! |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment 4 replies
-
The primary method for zero-copy transfer of GPU buffers into JAX is via the dlpack data structure, using the If dlpack is not an option, JAX also has support for |
Beta Was this translation helpful? Give feedback.
The primary method for zero-copy transfer of GPU buffers into JAX is via the dlpack data structure, using the
jax.from_dlpack
function.If dlpack is not an option, JAX also has support for
__cuda_array_interface__
.