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

arXiv: https://arxiv.org/abs/2510.09898 · HTML: https://mlsanity.com/paper/2510.09898/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 31 references that are in the map.
- First posted on arXiv 10 Oct 2025.

## Builds on (most cited first, 25 of 25)

- [CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation](https://mlsanity.com/paper/2102.04664/) (2021)
- [CodeSearchNet Challenge: Evaluating the State of Semantic Code Search](https://mlsanity.com/paper/1909.09436/) (2019)
- [GraphCodeBERT: Pre-training Code Representations with Data Flow](https://mlsanity.com/paper/2009.08366/) (2020)
- [CodeBLEU: a Method for Automatic Evaluation of Code Synthesis](https://mlsanity.com/paper/2009.10297/) (2020)
- [Unified Pre-training for Program Understanding and Generation](https://mlsanity.com/paper/2103.06333/) (2021)
- [UniXcoder: Unified Cross-Modal Pre-training for Code Representation](https://mlsanity.com/paper/2203.03850/) (2022)
- [CodeBERTScore: Evaluating Code Generation with Pretrained Models of Code](https://mlsanity.com/paper/2302.05527/) (2023)
- [HumanEval-XL: A Multilingual Code Generation Benchmark for Cross-lingual Natural Language Generalization](https://mlsanity.com/paper/2402.16694/) (2024)
- [Leveraging Automated Unit Tests for Unsupervised Code Translation](https://mlsanity.com/paper/2110.06773/) (2021)
- [DOBF: A Deobfuscation Pre-Training Objective for Programming Languages](https://mlsanity.com/paper/2102.07492/) (2021)
- [LLMs Assist NLP Researchers: Critique Paper (Meta-)Reviewing](https://mlsanity.com/paper/2406.16253/) (2024)
- [CodeTransOcean: A Comprehensive Multilingual Benchmark for Code Translation](https://mlsanity.com/paper/2310.04951/) (2023)
- [CodeUltraFeedback: An LLM-as-a-Judge Dataset for Aligning Large Language Models to Coding Preferences](https://mlsanity.com/paper/2403.09032/) (2024)
- [Exploring and Unleashing the Power of Large Language Models in Automated Code Translation](https://mlsanity.com/paper/2404.14646/) (2024)
- [ICE-Score: Instructing Large Language Models to Evaluate Code](https://mlsanity.com/paper/2304.14317/) (2023)
- [Does BLEU Score Work for Code Migration?](https://mlsanity.com/paper/1906.04903/) (2019)
- [CodeJudge: Evaluating Code Generation with Large Language Models](https://mlsanity.com/paper/2410.02184/) (2024)
- [Code Translation with Compiler Representations](https://mlsanity.com/paper/2207.03578/) (2022)
- [A Systematic Literature Review on Large Language Models for Automated Program Repair](https://mlsanity.com/paper/2405.01466/) (2024)
- [Large Language Models of Code Fail at Completing Code with Potential Bugs](https://mlsanity.com/paper/2306.03438/) (2023)
- [AlphaTrans: A Neuro-Symbolic Compositional Approach for Repository-Level Code Translation and Validation](https://mlsanity.com/paper/2410.24117/) (2024)
- [CodeRosetta: Pushing the Boundaries of Unsupervised Code Translation for Parallel Programming](https://mlsanity.com/paper/2410.20527/) (2024)
- [InterTrans: Leveraging Transitive Intermediate Translations to Enhance LLM-based Code Translation](https://mlsanity.com/paper/2411.01063/) (2024)
- [Open LLMs are Necessary for Current Private Adaptations and Outperform their Closed Alternatives](https://mlsanity.com/paper/2411.05818/) (2024)
- [Program Translation via Code Distillation](https://mlsanity.com/paper/2310.11476/) (2023)

## Abstract (arXiv)

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 T2J, 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. We then use T2J to improve PyTorch-to-JAX translation for the weak LLM gpt-4o-mini via in-context learning. Our evaluation shows that using T2J yields up to 20% improvement of our proposed metric T2J-CodeTrans-Score.