JAXbind: Bind any function to JAX

arXiv:2403.08847v1 Announce Type: cross
Abstract: JAX is widely used in machine learning and scientific computing, the latter of which often relies on existing high-performance code that we would ideally like to incorporate into JAX. Reimplementing the existing code in JAX is often impractical and the existing interface in JAX for binding custom code requires deep knowledge of JAX and its C++ backend. The goal of JAXbind is to drastically reduce the effort required to bind custom functions implemented in other …

abstract arxiv astro-ph.im code computing cs.lg function jax machine machine learning performance stat.co type