mirror of
https://github.com/ROCm/jax.git
synced 2025-04-17 20:36:05 +00:00
e1a68eee5e
This PR includes an end-to-end example project which demonstrates the use of the FFI. This complements [the FFI tutorial](https://jax.readthedocs.io/en/latest/ffi.html) 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).