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

Researchers introduce T2J, a benchmark dataset of PyTorch-to-JAX translation bugs paired with developer fixes, addressing the challenge of translating deep-learning code between frameworks. By training LLMs on this curated bug-fix data through in-context learning, they achieve up to 20% improvement in translation accuracy, demonstrating that domain-specific bug datasets can significantly enhance code generation reliability.

The translation of deep-learning code between frameworks represents a critical pain point in the AI development workflow. While large language models excel at translating between general-purpose programming languages, domain-specific code translation remains problematic because correctness depends on framework-specific APIs, execution semantics, and idiomatic patterns. PyTorch and JAX, despite both being popular deep-learning frameworks, have fundamentally different design philosophies and execution models, making naive translation error-prone.

The T2J benchmark addresses this gap by creating a structured learning resource. Rather than relying on theoretical understanding, the researchers collected actual bugs generated by LLMs, then hired developers to fix them. This ground-truth data enables LLMs to learn from concrete examples of what goes wrong and how corrections should be made. The 20% improvement in their proposed T2J-CodeTrans-Score metric suggests this approach is substantially more effective than baseline translation methods.

For the AI developer ecosystem, this work has practical implications. As teams increasingly adopt JAX for performance-critical applications, the ability to automatically translate existing PyTorch code reduces migration friction and speeds adoption. For LLM developers, this demonstrates that curated domain-specific datasets are more valuable than scale alone—even a weak model like gpt-4o-mini can achieve significant improvements through targeted in-context learning with relevant examples.

Future development likely involves expanding T2J beyond 20 kernels, exploring similar benchmarks for other framework pairs, and investigating whether this approach generalizes to other domain-specific translation tasks. The methodology itself—collecting real bugs and fixes rather than synthetic data—establishes a template for improving code generation in specialized domains.

• →T2J benchmark demonstrates that curated bug-fix datasets improve PyTorch-to-JAX translation accuracy by up to 20%
• →Domain-specific code translation requires framework-aware learning resources that go beyond general programming language translation
• →In-context learning with concrete examples outperforms baseline approaches even for smaller, less-capable LLM models
• →Reducing friction in deep-learning framework migration accelerates adoption of performance-optimized tools like JAX
• →The bug-fix collection methodology establishes a replicable pattern for improving code generation in specialized domains