JAX qfunctions
- Dominant language
- C
- Stars
- 265
- Forks
- 78
- Avg merge
- 2d 8h
- Merged PRs (30d)
- 11
Description
@jedbrown and I have been discussing the possibility of using JAX to write qfunctions, since it supports JIT compilation and automatic differentiation. I see several ways to go about this, and several potential roadblocks, so I'm opening this issue for discussion. First, we need to decide what sort of architecture we want -- here are a few options:
1. Implement an entire libCEED backend in JAX so you never have to leave Python (and write some simple C wrappers to call the Python functions from the C API, using the numpy C API). A major advantage of this approach would mean we don't need to worry about device memory handoff or code fusion, because everything would happen inside JAX, but a major disadvantage would mean that you would be tied to that backend for everything, even if a different backend for non-qfunction things would give better performance (i.e. if hand-written GPU code is faster than XLA-generated GPU code)
2. Only implement qfunctions themselves in JAX; i.e. write a C function to convert qfunction data into appropriate JAX `DeviceArray` instances. The major advantage of this approach would be that it's at least to some extent backend-independent (perhaps not the getting data into `DeviceArray`s part) and would require writing less Python code (i.e. not having to implement most libCEED functions in Python), but the biggest disadvantage would probably be that it's not necessarily easy to get the data into a JAX array on the device with no copying. The goal would be to avoid having to write C++ code that depends on XLA itself, since such code can really only be compiled in any reasonable manner by Bazel
If any of the libCEED devs have thoughts on this or are interested in working with me on implementing it, please let me know
Contributor guide
Assessment
This issue has not been assessed yet.