jax icon indicating copy to clipboard operation
jax copied to clipboard

[WIP] Add `ffi_call` tutorial

Open dfm opened this issue 1 year ago • 0 comments

Building on https://github.com/google/jax/pull/21925, this tutorial demonstrates the use of the FFI using ffi_call with a simple example. I don't think this should cover all of the most advanced use cases, but it should be sufficient for the most common examples. I think it would be useful to eventually replace the existing CUDA tutorial, but I'm not sure that it'll get there in the first draft.

As an added benefit, this also runs a simple test (akin to docs/cuda_custom_call) which actually executes using a tool chain that open source users would use in practice.

The current rendered version of the tutorial is here: https://jax--22095.org.readthedocs.build/en/22095/ffi/ffi.html

dfm avatar Jun 25 '24 18:06 dfm