Pith. sign in

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 →

arxiv 2510.09898 v2 pith:2S3DVO2Y submitted 2025-10-10 cs.LG cs.AI

classification cs.LGcs.AI
keywords PyTorch-to-JAXtranslationcodein-contextlearningbug-fixdatasetLLM-as-judgepromptaugmentationfunctionalcorrectnessBLEU
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper tries to establish that a weak LLM translating PyTorch code to JAX can be made dramatically more reliable by adding a curated dataset of previously observed translation bugs and their fixes into the prompt, with no fine-tuning. The authors built T2J, a benchmark that pairs 20 PyTorch kernels with developer-corrected JAX code and 163 error/fix entries, then used that data as in-context guidance for gpt-4o-mini. Their reported intrinsic results show the augmented prompt raises CodeBLEU from 0.19 to 0.29, raises a new LLM-judged functional-correctness score by about 0.95 points, and drops the number of human fix steps from 163 to 87. If those measurements hold, the core claim is that prompt-level bug context—not model retraining—can meaningfully close the gap for a low-cost model on a domain where the target library is underrepresented.

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.

Watch

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

Editorial extensions of the paper, not claims the author makes directly.

  • 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.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 5 minor

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)
  1. [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.
  2. [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.
  3. [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.
  4. [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)
  1. [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.
  2. [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.
  3. [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.
  4. [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.
  5. [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

1 steps flagged · score 6.0 of 10

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.

  1. 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 2 free parameters · 3 assumptions · 4 invented entities

The paper's claims rest on a narrow correctness definition, an unvalidated LLM-judge, an unverified LLM-generated reference for extrinsic evaluation, and a human fix-step count with no protocol. These are assumptions, not fitted constants; the central contribution is the small dataset, which is not openly shipped.

free parameters (2)
  • LLM sampling hyperparameters (temperature, top-p, max tokens) = unreported (ChatGPT-pro defaults)
    All LLM outputs and LLM-judge scores are generated through the ChatGPT-pro interface or OpenRouter API at default settings; no seed or sampling parameters are given, so results are not pinned down.
  • Human fix-step counting policy
    T2J FixCost Score is the number of fix steps; the granularity of a 'step' is defined by the two developers and is not captured in a repeatable protocol. This makes the central outcome measure (163->87) non-reproducible.
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.
    Section 4.2 defines verified correctness this narrow way; it does not require equivalence across untested inputs, so the human-fixed dataset may contain translations that are wrong on other inputs. This is load-bearing for both the dataset and the FixCost measurement.
  • ad hoc to paper JAX code generated by gpt-4o (without human verification) can serve as ground truth for extrinsic evaluation.
    Section 4.1.4 states extrinsic ground truth is 'JAX code generated by a costly LLM' because GitHub snippets lack test cases; no human verification is performed, so the reference itself may be incorrect. The same model is then used as judge (Appendix A.3).
  • domain assumption gpt-4o's LLM-as-judge scores reflect translation quality.
    The paper's main improvements are measured by T2J CodeTrans / Comparison Scores produced by gpt-4o. Table 3 shows these correlate only 0.04–0.29 (Pearson/Spearman) with the human FixCost, so this assumption is not supported by the paper's own data.
invented entities (4)
  • T2J CodeTrans Score
    purpose: LLM-as-judge metric (0-4) for usefulness and functional correctness of a JAX translation.
    No external validation; its correlation with human FixCost is weak (Table 3), and in extrinsic mode the judge is also the reference generator.
  • T2J FixCost Score
    purpose: Count of human fix steps needed to make a JAX translation correct.
    Depends on an unreported granularity of what counts as a fix step; no inter-annotator agreement is reported.
  • T2J Comparison Score
    purpose: LLM-as-judge pairwise preference between two translations given the source.
    A single gpt-4o prompt output; no reliability or bias analysis.
  • Fixed-bug dataset (T2J dataset)
    purpose: 163 bug/fix pairs from 20 TorchLeet problems, used as in-context examples.
    The dataset is a concrete resource but is not shipped with the paper; it cannot be inspected or reused from the text alone.

how reviews work

0 comments
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 reproduced from arXiv: 2510.09898 by the authors.

Figure 1
Figure 1. Example of PyTorch-to-JAX translation. (a) Input code; (b) Incorrect translation by 4o [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗
Figure 2
Figure 2. Overview Architecture of T2J . In Module One, a multi-step bug fixing process for the [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗
Figure 3
Figure 3. Standard Prompt for PyTorch-to-JAX code translation. [PITH_FULL_IMAGE:figures/full_fig_p012_3.png] view at source ↗
Figures from the paper (6 more)
Figure 4
Figure 4. Figure 4: Prompt for Augmenting to the weak LLM. The prompt in blue shows the immediate query after the prompt. {CODE}is the starting string of the code. A.3 EVALUATION PROMPT FOR T2J CODETRANS SCORE A.3.1 FUNCTIONAL CORRECTNESS See the prompt without reference at [PITH_FULL_IM…
Figure 5
Figure 5. Figure 5: Prompt for Scoring Functional Correctness by T2J [PITH_FULL_IMAGE:figures/full_fig_p013_5.png]
Figure 6
Figure 6. Figure 6: Prompt for Scoring Functional Correctness by T2J [PITH_FULL_IMAGE:figures/full_fig_p014_6.png]
Figure 7
Figure 7. Figure 7: Prompt for Scoring Usefulness by T2J CodeTrans Use NoRef [PITH_FULL_IMAGE:figures/full_fig_p015_7.png]
Figure 8
Figure 8. Figure 8: Prompt for Scoring Usefulness by T2J CodeTrans Use Ref You are an expert in PyTorch to JAX translation. I provide 3 inputs: 1 . PyTorch input code; 2. Translated Code Candidate A; 3. Translated Code Candidate B. Which candidate is a better translation result for this P…
Figure 9
Figure 9. Figure 9: Prompt for T2J Comparison Score. 16 [PITH_FULL_IMAGE:figures/full_fig_p016_9.png]

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

17 extracted references · 2 linked inside Pith

  1. [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...

  2. [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...

  3. [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...

  4. [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}...

  5. [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...

  6. [10]

    Assume the code was translated from PyTorch and should preserve its logic

  7. [11]

    Evaluate whether the JAX code appears complete, meaningful, and implementationally correct based on general expectations for such translations

  8. [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...

Show all 17 references
  1. [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...

  2. [16]

    Assume the PyTorch source code performs a well-defined functionality

  3. [17]

    Determine whether the JAX translated code snippet enables meaningful reuse or guidance toward equivalent implementation

  4. [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...

  5. [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}

  6. [20]

    Translated Code A: {TRANSLATE CODE A}

  7. [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...

  8. [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...

  9. [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,

Pith tools

Reviewed August 4, 2026 · model on record in the stance chip above.