JaxDecompiler: Redefining Gradient-Informed Software Design
Pierrick Pochelu
Among numerical libraries capable of computing gradient descent optimization,
JAX stands out by offering more features, accelerated by an intermediate
representation known as Jaxpr language. However, editing the Jaxpr code is not
directly possible. This article introduces JaxDecompiler, a tool that
transforms any JAX function into an editable Python code, especially useful for
editing the JAX function generated by the gradient function. JaxDecompiler
simplifies the processes of reverse engineering, understanding, customizing,
and interoperability of software developed by JAX. We highlight its
capabilities, emphasize its practical applications especially in deep learning
and more generally gradient-informed software, and demonstrate that the
decompiled code speed performance is similar to the original.