Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

JaxDecompiler: Redefining Gradient-Informed Software Design

Pierrick Pochelu

TL;DR

JaxDecompiler simplifies the processes of reverse engineering, understanding, customizing, and interoperability of software developed by JAX, especially useful for editing the JAX function generated by the gradient function.

Abstract

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.

JaxDecompiler: Redefining Gradient-Informed Software Design

TL;DR

JaxDecompiler simplifies the processes of reverse engineering, understanding, customizing, and interoperability of software developed by JAX, especially useful for editing the JAX function generated by the gradient function.

Abstract

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.
Paper Structure (12 sections, 1 figure, 2 tables)

This paper contains 12 sections, 1 figure, 2 tables.

Table of Contents

  1. Introduction
  2. Step-by-step use cases
  3. Step 1: JaxDecompiler input
  4. Step 2: Utilizing JaxDecompiler API
  5. Step 3: Consuming JaxDecompiler output
  6. JaxDecompiler design
  7. Overview
  8. Tokenizer
  9. Line Translator
  10. Import Set
  11. Performance of the decompiled code
  12. Conclusion

Figures (1)

  • Figure 1: Design of the JaxDecompiler. Edge represents data flow and the box the processing.