REVIEW 4 major objections 5 minor 17 references
Learning Bug Context for PyTorch-to-JAX Translation with LLMs
T0 review · 4 major / 5 minor · reviewed 2026-08-04 · deepseek-v4-flash
Pith's one-line read In-context bug-fix examples from past translation errors improve a weak LLM's PyTorch-to-JAX output, cutting human repair effort by nearly half.
desk verdict A small but real human-annotated PyTorch-to-JAX bug-fix dataset; the prompt-augmentation story is plausible, but the LLM-judge metrics are too weakly validated to carry the headline. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The fixed-bug dataset: a JSON array of 163 error/fix entries derived from 20 TorchLeet PyTorch programs translated by gpt-4o-mini and repaired by two professional developers to runnable, functionally equivalent JAX. Each entry records the buggy JAX code, the error message, a textual fix instruction, and the corrected code. The augmented prompt pastes this dataset alongside the new PyTorch snippet, so the weak LLM can pattern-match against previously seen error categories (training loops, rng/PRNG misuse, model-definition bugs) rather than translating from scratch.
What would settle it
Run a small human study: take the 20 intrinsic samples, have independent developers (not the paper's) rank baseline vs T2J outputs for correctness, and compare their ranking with the T2J FixCost scores and with gpt-4o's CodeTrans scores. If human judgments disagree with gpt-4o on a substantial fraction of samples, the headline gains would not reflect true translation quality.
Extended reading notes
Core claim
The central claim is that a fixed-bug dataset created by human developers repairing LLM-generated JAX code can be packaged as structured JSON context and injected into the prompt of a cheap LLM, steering it past recurring PyTorch-to-JAX pitfalls. On the 20-sample intrinsic set, the T2J pipeline reports a 10% relative CodeBLEU gain (0.19 to 0.29), a 0.8-point usefulness gain and 0.95-point functional-correctness gain on the 0-4 scale, a 100% win-rate on the comparison metric, and a reduction in total human fix steps from 163 to 87. On a 100-sample GitHub-derived extrinsic set, the pipeline reports gains on the LLM-judged metrics while CodeBLEU actually drops slightly, which the authors attrib
Load-bearing premise
The reported quality gains rest on the validity of gpt-4o's judge scores (and, in the extrinsic setting, gpt-4o's untested code as ground truth) at a time when those same scores correlate only 0.04–0.29 with human-measured fix cost.
Editorial extensions
If this is right
- If reproduced on larger sets, prompt-level bug context offers a low-cost path to improving code translation without model retraining, useful for other under-resourced library pairs.
- The FixCost drop (163 to 87 steps) suggests downstream human debugging effort can be roughly halved by teaching the model common failure modes in-context.
- The proposed T2J CodeTrans and Comparison metrics give a template for evaluating cross-library translation when reference code is unavailable, as in repository-level extrinsic settings.
- The model's running-time improvement (about 2.5x faster than baseline) hints that fewer fix iterations also yield more idiomatic, vectorizable JAX, though the paper does not deeply analyze why.
- The weak correlations between the LLM-judged metrics and human FixCost (Pearson 0.04–0.29, Table 3) imply the headline gains may partly reflect judge preferences rather than measurable human repair burden.
Reading between the lines
- A direct test of the framework's validity would be to have a human judge, blind to method, rank baseline vs T2J outputs on a held-out set; the paper's own correlation table suggests LLM-as-judge scores may not align tightly with human fixing effort.
- If the approach generalizes, the same error-fix JSON structure could be transferred to other domain-specific translations (e.g., NumPy-to-JAX, PyTorch-to-TensorFlow), where similar underrepresentation in training data is the bottleneck.
- An extension the paper leaves implicit is weighting fix steps by difficulty; treating every fix as one unit likely underestimates the practical gain when complex fixes (e.g., jit/grad rewrites) are avoided.
- The extrinsic evaluation's reliance on untested gpt-4o output as ground truth is a notable pillar; a future benchmark with executable test cases for GitHub-sourced code would clarify whether T2J's extrinsic gains are genuine.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper introduces T2J, a framework that augments prompts to a weak LLM (GPT-4o-mini) with a curated dataset of human-verified PyTorch-to-JAX translation bugs and fixes, with the goal of improving translation quality. The authors construct the fixed-bug dataset from 20 TorchLeet problems, hire two developers to repair GPT-4o-mini's JAX drafts, and then use the resulting bug/fix pairs as in-context examples. They also propose three new evaluation metrics—T2J CodeTrans Score, T2J FixCost Score, and T2J Comparison Score—and report intrinsic results on the 20 TorchLeet examples and extrinsic results on 100 GitHub-sourced PyTorch snippets. The reported improvements include CodeBLEU rising from 0.19 to 0.29, human fix cost dropping from 163 to 87 steps, and large gains on the proposed LLM-judged metrics.
Significance. If the claims hold, the fixed-bug dataset and the prompt-augmentation idea are a useful contribution to the underexplored PyTorch-to-JAX translation problem, especially for low-cost LLMs. The paper's human-curated bug/fix pairs provide a concrete resource, and the measured reduction in human fixing effort (163 to 87 steps) is a tangible, human-grounded result. However, the significance is currently limited by the evaluation's validation gaps: the proposed metrics correlate only weakly with human fix cost (Table 3), and the extrinsic evaluation uses GPT-4o output as both the reference and the judge. These issues directly affect the credibility of the headline gains.
major comments (4)
- [Section 4.1.4, Section 5.2.1, Figure 2 caption] The extrinsic evaluation is circular: the 'ground-truth' JAX code is generated by GPT-4o (Section 4.1.4), and the T2J CodeTrans and Comparison scores are also produced by GPT-4o prompts (Appendix A.3). Thus, substantial reported gains in the extrinsic setting—e.g., +1.2 usefulness and +0.6 functional correctness in Section 5.2.1—may reflect the judge's preference for outputs similar to its own reference rather than actual translation quality. The figure caption itself states that human bug-fixing was not performed on the GitHub corpus because repository-level snippets lack test cases. Please provide an independent validation of the judge on a human-annotated/extrinsic sample, or restrict the claims to the intrinsic, human-verified FixCost results.
- [Section 5.2.2, Table 3] Table 3 shows that the proposed T2J CodeTrans metrics correlate with the human-measured T2J FixCost Score at only 0.04–0.29 (Pearson) and 0.07–0.29 (Spearman). The authors acknowledge these are weak correlations. Since the T2J CodeTrans and Comparison scores are the basis for most of the reported improvements in Section 5.2.1, the weak correlation means these metrics are not established as valid measures of translation quality. The paper should either calibrate the proposed metrics against human/execution judgments, provide evidence that the LLM judge is reliable, or substantially temper the claims that rely on these unvalidated metrics.
- [Section 4.2, Eq. (1), Section 6] The FixCost protocol is underspecified. The paper defines T2J FixCost Score as the number of fix steps but never defines what constitutes one 'fix step' (e.g., a single edit, a semantic unit, or an error-message-guided change), nor does it report per-sample counts, inter-rater reliability, or whether the developers were blinded to the experimental condition. The Limitations section (Section 6) admits that fixes vary in difficulty and that only step counts are used. With n=20 and no variance or protocol details, the headline 163→87 reduction cannot be assessed for robustness or statistical significance.
- [Section 4.4, Table 2] The intrinsic evaluation is based on only 20 samples, with no confidence intervals or significance tests. The text in Section 4.4 says cross-validation uses 19 samples to translate the remaining sample, but Table 2 reports only aggregate values. In addition, the CodeBLEU improvement is described in Section 5.2.1 as 'a 10% relative gain' for 0.19→0.29; the relative gain is actually about 53%, while the absolute gain is 0.10. Please clarify the intended statement and report variance (e.g., bootstrap confidence intervals) for all intrinsic metrics.
minor comments (5)
- [Abstract and Section 1] There are several typos and grammar issues, e.g., 'a in-context code learning', 'Pytorch' and 'JAX' inconsistent capitalization, and 'thich' in Appendix A.2. A careful proofread is recommended.
- [Figure 1] In the motivation example, the text says the error is in 'Line 8' of Figure 1(b), but the incorrect shape appears in line 9 (`W = random.normal(w_key, (10, 28*28))`). Please correct the reference.
- [References] Several references use placeholder access dates ('Accessed: YYYY-MM-DD'), and some URLs are missing publication year or venue details. Please update the bibliography.
- [Table 2 / Section 5.2.1] The extrinsic CodeBLEU result (Baseline 0.41 vs T2J 0.38, a 3% decrease) is described as the baseline 'outperformed T2J by 3%'—this should be stated as an absolute or relative difference explicitly, since the metric direction is not obvious from the percentage alone.
- [Section 5.1] The experimental setup reports using ChatGPT-pro's interface for code generation and API/OpenRouter for evaluation. The exact model versions and temperature settings are not reported; please provide these details to make the evaluation reproducible.
Circularity Check
Extrinsic evaluation is self-referential: gpt-4o supplies the ground-truth reference and also judges candidates against it, so the extrinsic CodeTrans/Comparison improvements are not independent; intrinsic FixCost/CodeBLEU results are independently grounded in human fixes.
-
self definitional
[Section 4.1.4 (Types of Evaluation), Section 4.3 (Evaluation Framework), Section 4.4 (Extrinsic Evaluation), Table 2]
"for intrinsic evaluation, the ground-truth JAX code is obtained through human verification and bug-fixing, whereas for extrinsic evaluation, we rely on JAX code generated by a costly LLM for the GitHub dataset. ... Second, the extrinsic evaluation considered the output of costly LLMs, i.e. JAX translated code by this process) as the ground-truth data point. ... we use the GPT-4o model to evaluate the quality of translated code by usefulness and functional correctness."
In the extrinsic setting the reference/ground-truth JAX code is gpt-4o's own translation of the GitHub PyTorch snippet, and the CodeTrans and Comparison scores are produced by the same gpt-4o model (Appendix A.3). The reported extrinsic gains (e.g., CodeTrans Use Ref 2.74 to 2.94, Func Ref 2.43 to 3.02, Comparison Score 0.18 to 0.82) therefore measure agreement/preference of gpt-4o with gpt-4o-generated references, not verified semantic equivalence. This is a self-referential evaluation loop: the model that defines 'correct' also determines 'quality'. The intrinsic FixCost/CodeBLEU results remain independent because they are grounded in human bug-fixing, so the circularity is partial.
full rationale
The derivation chain for the central claim is a dataset-construction plus in-context-learning pipeline, not a mathematical derivation. The intrinsic evaluation is genuinely non-circular: it uses leave-one-out cross-validation over 20 TorchLeet problems, and the ground truth is human-fixed JAX code; the FixCost improvement (163 to 87 fixing steps) and CodeBLEU gain (0.19 to 0.29) are measured against this independent human-verified standard. No load-bearing self-citation or imported uniqueness theorem is used. The serious issue is the extrinsic evaluation: the paper explicitly defines gpt-4o output as the ground truth (Sections 4.1.4 and 4.4) while also using gpt-4o as the LLM judge for T2J CodeTrans and Comparison scores (Section 4.3). Thus the extrinsic CodeTrans/Comparison improvements in Table 2 are self-referential and cannot independently establish translation quality. This is a partial circularity rather than a complete one, because the intrinsic FixCost and CodeBLEU columns do not depend on that loop. The paper's own Table 3 and Section 5.2.2 additionally report weak correlations (<=0.29) between the LLM-judge metrics and human FixCost, but that is a validity concern rather than an additional derivation-level circle.
Assumptions & free parameters
free parameters (2)
- LLM sampling hyperparameters (temperature, top-p, max tokens) =
unreported (ChatGPT-pro defaults)
- Human fix-step counting policy
assumptions (3)
- domain assumption A translated JAX snippet is 'correct' iff it compiles, runs, and returns the same output as the PyTorch code on the given test case.
- ad hoc to paper JAX code generated by gpt-4o (without human verification) can serve as ground truth for extrinsic evaluation.
- domain assumption gpt-4o's LLM-as-judge scores reflect translation quality.
invented entities (4)
-
T2J CodeTrans Score
-
T2J FixCost Score
-
T2J Comparison Score
-
Fixed-bug dataset (T2J dataset)
Cite this review
Pith. "Pith review of Learning Bug Context for PyTorch-to-JAX Translation with LLMs." pith.science (2026). https://pith.science/paper/2S3DVO2Y
@misc{pith2026251009898,
author = {Pith},
title = {Pith review of: Learning Bug Context for PyTorch-to-JAX Translation with LLMs},
year = {2026},
howpublished = {\url{https://pith.science/paper/2S3DVO2Y}},
note = {Machine review of arXiv:2510.09898}
}
read the original abstract
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.
Figures
Figures from the paper (6 more)
Reference graph
Works this paper leans on
-
[2]
Error Code
A JSON file that contains a dataset of common errors in PyTorch-to-JAX translation by Weak LLM 4o-mini. Each data point contains the following fields: -Example id: ID of the source code. -Input Code: Source code in Pytorch. -LLM weak output: JAX translated code ofInput Codeusing a weak LLM (4o-mini). LLM fix output: Fixed JAX code fromLLM weak outputby th...
-
[3]
Ngoc Tran, Hieu Tran, Son Nguyen, Hoan Nguyen, and Tien N
URLhttps://arxiv.org/abs/2410.02184. Ngoc Tran, Hieu Tran, Son Nguyen, Hoan Nguyen, and Tien N. Nguyen. Does bleu score work for code migration? InProceedings of the 27th International Conference on Program Com- prehension, ICPC ’19, pp. 165–176. IEEE Press, 2019. doi: 10.1109/ICPC.2019.00034. URL https://doi.org/10.1109/ICPC.2019.00034. Yue Wang, Hung Le...
arXiv 2019
-
[4]
Shuyan Zhou, Uri Alon, Sumit Agarwal, and Graham Neubig
URLhttps://arxiv.org/abs/2405.01466. Shuyan Zhou, Uri Alon, Sumit Agarwal, and Graham Neubig. Codebertscore: Evaluating code generation with pretrained models of code, 2023. URLhttps://arxiv.org/abs/2302. 05527. Ming Zhu, Mohimenul Karim, Ismini Lourentzou, and Daphne Yao. Semi-supervised code trans- lation overcoming the scarcity of parallel code data. I...
arXiv 2023
-
[6]
Your task is to reason and get the output JAX code from these above inputs
The data.csv file thich stored possible input when running some examples in the JSON file. Your task is to reason and get the output JAX code from these above inputs. Please note that you can learn the process of error fixing in Torch-to-JAX translation in 2) JSON file. Now I will give you a set of input in the next query. Input Source Code Snippet:{CODE}...
-
[9]
Assign a score for functional correctness on a scale from 0 to 4. Input Source Code in PyTorch: {SOURCE CODE} Translated JAX Code Snippet: {TRANSLATED CODE} Evaluation Form: Functional Correctness (scores ONLY): Figure 5: Prompt for Scoring Functional Correctness by T2J CodeTrans Func NoRef Table 6: Error categories and their counts in Human Bug Fixing da...
-
[10]
Assume the code was translated from PyTorch and should preserve its logic
-
[11]
Evaluate whether the JAX code appears complete, meaningful, and implementationally correct based on general expectations for such translations
-
[12]
Assign a score for functional correctness on a scale from 0 to 4. Input Source Code in PyTorch: {SOURCE CODE} Translated JAX Code Snippet: {TRANSLATED CODE} Reference JAX Code Snippet: {REFERENCE} Evaluation Form: Functional Correctness (scores ONLY): Figure 6: Prompt for Scoring Functional Correctness by T2J CodeTrans Func Ref Table 7: Error subcategorie...
2025
Show all 17 references
-
[15]
Assign a score for usefulness from 0 to 4. Input Source Code in PyTorch: {SOURCE CODE} Translated JAX Code Snippet: {TRANSLATED CODE} Evaluation Form: Usefulness (scores ONLY): Figure 7: Prompt for Scoring Usefulness by T2J CodeTrans Use NoRef Table 8: Error subcategories unde...
-
[16]
Assume the PyTorch source code performs a well-defined functionality
-
[17]
Determine whether the JAX translated code snippet enables meaningful reuse or guidance toward equivalent implementation
-
[18]
Assign a score for usefulness from 0 to 4. Input Source Code in PyTorch: {SOURCE CODE} Translated JAX Code Snippet: {TRANSLATED CODE} Reference JAX Code Snippet: {REFERENCE} Evaluation Form: Usefulness (scores ONLY): Figure 8: Prompt for Scoring Usefulness by T2J CodeTrans Use...
-
[19]
Translated Code Candidate B
Translated Code Candidate A; 3. Translated Code Candidate B. Which candidate is a better translation result for this Pytorch code. Input Pytorch code: {CODE}
-
[20]
Translated Code A: {TRANSLATE CODE A}
-
[21]
Figure 9: Prompt for T2J Comparison Score
Translated Code B: {TRANSLATE CODE B} Please also provide the reason why you consider a candidate better than the other translated code candidate. Figure 9: Prompt for T2J Comparison Score. 16 Table 9: Example results comparing PyTorch, ground truth, baseline, and T2J (T2J) ou...
-
[2020]
Baptiste Roziere, Marie-Anne Lachaux, Marc Szafraniec, and Guillaume Lample
URLhttps://arxiv.org/abs/2009.10297. Baptiste Roziere, Marie-Anne Lachaux, Marc Szafraniec, and Guillaume Lample. Dobf: A deob- fuscation pre-training objective for programming languages, 2021. URLhttps://arxiv. org/abs/2102.07492. Baptiste Roziere, Jie M. Zhang, Francois Char...
2009 arXiv
-
[2024]
Weixi Tong and Tianyi Zhang
URLhttps://openreview.net/forum?id=V6hrg4O9gg. Weixi Tong and Tianyi Zhang. Codejudge: Evaluating code generation with large language models,
Reviewed August 4, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.