JaxDecompiler: Redefining Gradient-Informed Software Design

Pierrick Pochelu
3/14/2024
cs.PLcs.LGcs.SE

No official code found. Best implementation: PierrickPochelu/JaxDecompiler (16 stars, Apache-2.0, JAX, last commit 2025-04-22).

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.

View on arXivView PDF

Code Implementations(1)

PierrickPochelu/JaxDecompilerCommunity100%
JAXApache-2.0

Jax Decompiler

161Shell, PythonDec 3, 20221 years agoApache-2.0
compilerjaxpython

Similar Papers

MPAX: Mathematical Programming in JAX

79% match1 repo1632024

JAXbind: Bind any function to JAX

79% match1 repo822024

JaxPruner: A concise library for sparsity research

78% match1 repo2362023

DaCe AD: Unifying High-Performance Automatic Differentiation for Machine Learning and Scientific Computing

76% match2025

Learning to Optimize by Differentiable Programming

76% match2026

Enhancing Deep Learning with Optimized Gradient Descent: Bridging Numerical Methods and Neural Network Training

75% match2024

Cite this paper

@article{pochelu2024jaxdecompiler,
  title  = {JaxDecompiler: Redefining Gradient-Informed Software Design},
  author = {Pierrick Pochelu},
  year   = {2024},
  eprint = {2403.10571},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2403.10571v1}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2403.10571.svg)](https://codeofpaper.com/paper/2403.10571)

Discussion