Learning Bug Context for PyTorch-to-JAX Translation with LLMs

\begin{abstract} % Code Translation Large language models (LLMs) have shown strong performance on code translation between widely used programming languages. However, translation becomes much less reliable for domain-specific code, where correctness depends on framework-specific APIs and execution semantics. One example is translating deep-learning code from PyTorch to JAX, where LLM outputs often contain subtle bugs or non-idiomatic usage that prevents execution or changes behavior. Prior work suggests that curated bug–fix data from LLM-generated code can help improve code generation quality, but such resources are still limited for PyTorch-to-JAX translation. In this work, we introduce \ourtoolnospace, a benchmark of LLM translation bugs paired with developer-written fixes for PyTorch-to-JAX code. We start from 20 kernels in the TorchLeet dataset, translate them to JAX using the weak LLM gpt-4o-mini, and hire software developers to debug and repair the generated JAX implementations. In total, \ourtool contains over 160 bug–solution pairs collected from real LLM outputs. We then use \ourtool to improve PyTorch-to-JAX translation for the weak LLM gpt-4o-mini via in-context learning. Our evaluation shows that using \ourtool yields up to 20\% improvement of our proposed metric $\ourtoolnospace\_CodeTrans\_Score$. We release our replication package at here\footnote{https://tinyurl.com/4ehky8hb} % , demonstrating the value of developer-curated bug–fix supervision for more reliable framework-level code translation. % Large Language Models have many success in general code translation between well-known programming languages. However, in the domain-specific languages code translation, there are risks of getting low quality translation results due to the lack of models' knowledge about the source/ target languages. Previous works show that collecting a dataset of bugs pattern from Large Language Models (LLM) generated code are needed to improve the quality of code generation on general programming languages. In this work, we select the research problem as code translation from PyTorch to JAX code. To overcome the challenge of lacking of bug-solution of LLM generated code for PyTorch-to-JAX translation, we collect a benchmark of LLM bugs and hired software developers to conduct fixing solution for each bugs appeared in LLM generated JAX code over 20 kernels of the well-known TorchLeet dataset in PyTorch. In total, we collect over 400 bugs and bugs' solutions from JAX generated code of two weak LLMs: gpt-4o-mini and QwenCoder-2.5-7B-IT. We use this benchmark, T2J-bench, to improve the quality in code translation of numerous open LLMs by two directions: in-context learning and fine-tuning with T2J-bench. Evaluation shows that T2J-bench can achieve up to XX\% of improvement in terms of ICE-score, which shows the potential of this benchmark in improving PyTorch-to-JAX code translation. \end{abstract}