jax icon indicating copy to clipboard operation
jax copied to clipboard

Add FFI example project and test on CI

Open dfm opened this issue 1 year ago • 0 comments

This PR includes an end-to-end example project which demonstrates the use of the FFI. This complements the FFI tutorial by putting all of the code in one place, as well as demonstrating how FFI extensions can be packaged. Alongside the example project, I have also added a new GitHub Actions workflow to test the example as part of CI.

For now, the tests only run on CPU, but once we have GPU runners for GitHub Actions (soon!), I plan on migrating the custom call examples from docs/gpu_ops and docs/cuda_custom_call into this test case.

Similarly, I wanted to start small and this example project only includes exactly the same functions as the tutorial for now, but I think this could be a good place to showcase more advanced examples (including custom calls with state).

dfm avatar Sep 03 '24 18:09 dfm