[JOSS 6532] Clarify related work

[wsmoses]

The package certainly has plenty of value for making it easy to declare custom derivatives of arbitrary python functions (including those which bind to other languages).

However, I do believe that it does not sufficiently address related work. For example "To the best of our knowledge no other code currently exists for connecting generic functions and both of their Jacobian products to JAX." This feels wrong and should be reworded. JaX certainly has the ability to bind custom python code to itself, alongside its forward and reverse mode derivatives. In fact those are the mechanisms used here. Similarly, Enzyme-JaX permits binding arbitrary programs, not just C++ as external code into JaX. In particular it can bind arbitrary LLVM/MLIR, including C++, and code JaX itself. It also can automatically generate the derivative of these, but similarly you can import a custom derivative as well.

I'm not quite sure how I would rephrase this offhand. Something like "connecting generic Python functions " feels closer, but not sufficient, given for example JaX itself and Enzyme can also provide this (albeit not in the same way). Perhaps it would be better to focus on usability in the statement of need. Something like "to our knowledge there exists no package for easily binding existing Python code into JaX" or perhaps "for binding generic Python code into JaX without knowledge of C++ or LLVM". The first option makes clear that it is easier than either utility at the moment (e.g. having to do the corresponding registration via C++ in JaX itself, or compiling the python into LLVM/MLIR for Enzyme-JaX's current frontend). The second option relies on the fact that JaX and Enzyme-JaX can only define derivatives for JaX-compatible python programs presently (e.g. those without non-JaX FFI) if one wants to work solely in Python.

There are plenty of other options here as well.

Also very minor for Enzyme-JaX you should cite Moses, William S. and Zinenko, Oleksandr not "EnzymeAD"

ref: https://github.com/openjournals/joss-reviews/issues/6532

[hawkinsp]

I'd second this. I think the key thing JAXbind is offering is better ergonomics than other competing approaches to extending JAX, as things stand. That's worth highlighting and something to be proud of.

[roth-jakob]

Thank you for your comments! I fully agree with this point. We are aware that there are multiple options for binding custom code to JAX. With PR #25 , we follow your suggestion, emphasizing that JAXbind provides an easy-to-use interface.

And sorry for messing up with the Enzyme-JAX author field in the bib file.

[roth-jakob]

Thanks again for your feedback! Please let us know in case you believe further clarification is needed.

[wsmoses]

Hm I'd lke to reopen this. I'm not sure my concerns about the representation of Enzyme-JaX has been addressed yet.

Per above "Similarly, Enzyme-JaX permits binding arbitrary programs, not just C++ as external code into JaX. In particular it can bind arbitrary LLVM/MLIR, including C++, and code JaX itself. It also can automatically generate the derivative of these, but similarly you can import a custom derivative as well."

I think perhaps changing "one to differentiate a C++ function" to "one to import and optionally differentiated LLVM/MLIR-based programs, like C++" would help. Regarding "it enforces the use of Enzyme for deriving derivatives and does not allow for connecting arbitrary code to JAX." perhaps striking "enforces the use of Enzyme for defining derivatives" and changing "does not allow for connecting arbitrary code to JAX" to "does not allow connecting arbitrary Python code to JaX".

[danielskatz]

For JOSS review https://github.com/openjournals/joss-reviews/issues/6532

[Edenhofer]

Yes, the wording isn't optimal yet. Maybe a much simpler statement as proposed in #29 just saying what Enzyme-JAX does is sufficient. What do you think?

[wsmoses]

Yup that sounds good to me!