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

Technical Summary: Learning Bug Context for PyTorch-to-JAX Translation with LLMs

Problem Statement
While Large Language Models (LLMs) demonstrate strong performance in translating code between general-purpose programming languages, their reliability diminishes significantly when handling domain-specific code, such as deep learning frameworks. Specifically, translating code from PyTorch to JAX presents unique challenges due to fundamental differences in execution semantics, including dynamic graph execution, just-in-time (JIT) compilation, automatic differentiation approaches, and vectorization flows. Prior work, such as UniTrans, attempts to improve translation via iterative correction using test inputs or compilation errors. However, these methods often fail to rectify translations effectively, particularly when relying on smaller or "weak" LLMs (e.g., those with fewer than 10 billion parameters), resulting in non-functional code or subtle behavioral changes that deviate from the original intent.

Methodology: The T2J Framework
The authors introduce T2J, a prompt augmentation framework designed to enhance LLM-based PyTorch-to-JAX translation by leveraging curated datasets of bugs and their fixes. The framework operates through several key modules:

1. Data Construction (Fixed-Bug Dataset):

   • The authors start with 20 PyTorch code snippets from the TorchLeet dataset (problem-solving code).
   • These snippets are translated into JAX using a "weak" LLM (gpt-4o-mini).
   • Professional software developers are hired to debug and repair the generated JAX implementations. This process involves multiple rounds of manual verification to ensure the JAX code is compilable, runnable, and produces outputs equivalent to the original PyTorch code.
   • This results in a curated dataset of 163 bug–solution pairs, documenting prevalent translation errors (e.g., API mismatches, gradient discrepancies) and their corresponding fixes.

2. In-Context Learning via Prompt Augmentation:

   • Instead of fine-tuning the model, T2J employs in-context learning.
   • The translation prompt is augmented with structured guidance derived from the fixed-bug dataset. Specifically, the prompt includes a JSON-formatted list of common error patterns and their solutions as context.
   • The "weak" LLM (gpt-4o-mini) uses this augmented context to generate corrected JAX code without requiring expensive retraining.

3. Evaluation Framework:

   • Datasets: Evaluation is conducted on two datasets: an Intrinsic set (20 TorchLeet snippets with human-verified ground truth) and an Extrinsic set (100 PyTorch snippets from GitHub, where ground truth is approximated by a "strong" LLM, gpt-4o, followed by human verification).
   • Metrics: The paper proposes three novel metrics alongside the traditional CodeBLEU:
     • T2J CodeTrans Score: An LLM-as-a-judge metric evaluating Usefulness and Functional Correctness (with and without reference code), inspired by ICE-score.
     • T2J FixCost Score: Quantifies the effort required to correct code by counting the number of human fix steps.
     • T2J Comparison Score: A binary judgment by an LLM comparing two translation outputs to determine which is superior.

Key Contributions

1. T2J Dataset: The creation of the first fixed-bug dataset specifically for PyTorch-to-JAX translations, containing detailed annotations of error patterns and fixes to facilitate reliability improvements in LLM code translation.
2. T2J Framework: A novel prompt augmentation technique that bridges domain-specific gaps by integrating structured bug-fix context, offering a scalable approach for cross-ecosystem migrations without the computational cost of fine-tuning.
3. Empirical Validation: Rigorous evaluation demonstrating that the proposed framework significantly enhances translation quality for weak LLMs.

Results

• Intrinsic Evaluation: Using T2J, the weak LLM (gpt-4o-mini) achieved a 20% improvement in the proposed T2J CodeTrans Score (specifically in functional correctness and usefulness) compared to the baseline standard prompt. The T2J Comparison Score indicated that 100% of T2J-generated translations were judged superior to the baseline.
• Efficiency: The framework reduced the human effort required for debugging. The T2J FixCost Score dropped from 163 steps (baseline) to 87 steps, representing roughly a 50% reduction in correction effort.
• Performance: The corrected code generated by T2J ran approximately 2.5 times faster than the baseline output in execution time tests.
• Extrinsic Evaluation: On the GitHub dataset, T2J showed improvements in functional correctness and usefulness (up to 1.2 points), though CodeBLEU scores were slightly lower than the baseline, highlighting the limitations of traditional metrics for this specific domain.
• Correlation: The study found that the proposed T2J CodeTrans metrics showed the strongest correlation with human fixing costs compared to traditional metrics like CodeBLEU.

Significance and Claims
The paper claims that T2J provides a cost-effective and scalable solution for improving cross-framework code translation. By curating a "learning bug context" dataset and utilizing in-context learning, the authors demonstrate that even weaker, low-cost LLMs can achieve high-quality translations that approach the performance of stronger models, provided they are guided by structured examples of common errors and their resolutions. The authors position this work as a step toward modernizing legacy deep learning systems and facilitating migration between frameworks, while acknowledging limitations such as the current focus on problem-solving datasets and the lack of application to open-source LLMs due to budget constraints.