JAXbind: Bind any function to JAX

JAX在机器学习和科学计算中得到了广泛应用，后者通常依赖于现有的高性能代码，我们希望将其理想地纳入JAX中。重新实现JAX中的现有代码通常是不切实际的，而JAX中用于绑定自定义代码的现有接口需要深入了解JAX及其C++后端。JAXbind的目标是大大减少将其他编程语言中实现的自定义函数绑定到JAX所需的工作量。具体而言，JAXbind提供了一个易于使用的Python接口，用于定义所谓的JAX原语，支持任意JAX变换。

JAXbind试图解决将现有的高性能代码绑定到JAX的问题。

JAXbind提供了一个易于使用的Python接口，用于定义自定义JAX原语，支持任意JAX变换。

JAXbind可以显著降低将其他编程语言中实现的自定义函数绑定到JAX所需的工作量。该论文还介绍了JAXbind的实现细节和性能评估结果。

与JAXbind相关的研究尚未发现。