用智能代理自动把PyTorch模型转成JAX，准确率达91%。
Agentic Framework for Deep Learning workload migration via In-Context Learning
- 用上下文学习+执行真值校验构建自主调试循环。
- 在真实张量状态基础上生成测试，实现91%数值等价。
- 适合需要跨框架迁移深度学习模型的研究者。
将深度学习模型从PyTorch的灵活面向对象设计迁移到JAX的函数式无状态架构，通常需手动完成且易出错。自动化迁移困难，因大语言模型（LLM）难以精确对齐动态API，且对精细操作易犯错。本文提出一个完全自治的系统，结合上下文学习（ICL）与基于真值的自调试机制。首先，构建包含典型JAX风格和测试用例生成规范的ICL上下文作为严格参考。其次，不依赖LLM推断数学结果，而是运行原始PyTorch模块获取实际动态张量状态，形成不可更改的执行真值。随后通过自主智能体循环基于真值数据合成测试，反复执行并回传错误堆栈供LLM自我修正。消融实验表明，结合ICL参考、真值锚定与自调试的方法显著优于纯指令型与基础智能体基线。该方法计算开销低，轻量级流程在神经模块上实现91%数值等价（基线仅9%，指令+自调试为27%）。已在SAM、T5、Code Whisper等多个前沿模型上验证，保持高数值一致性。
原文摘要 · Abstract (English)
Translating deep learning models from PyTorch's flexible, object-oriented design to JAX's functional, stateless setup is usually a manual and error-prone task. Automated migration is challenging because Large Language Models (LLMs) struggle with strict and dynamic API alignment and are prone to mistakes for exacting operations. We propose a fully autonomous system that combines In-Context Learning (ICL) with oracle-driven self-debugging. First, we curated an ICL context that serves as a strict reference for idiomatic JAX styling and test case generation. Second, instead of depending on the LLM to deduce mathematical outputs, we run the source PyTorch modules to get their actual dynamic tensor states. This creates an unchangeable execution oracle. We then use an autonomous agentic loop to synthesize tests based on the oracle data. The test cases are executed repeatedly, and the traceback is sent back to the LLM for self-correction. Ablations show that combining ICL references with oracle grounding and self-debugging greatly outperforms pure instructional and basic agentic baselines. This improvement does not add an excessive computational overhead. Our lightweight pipeline achieves 91% numerical equivalence (compared to baseline: 9%, instruction + self-debugging: 27%) on neural modules, providing a highly reliable, scalable blueprint for cross-framework migration. This has been validated across several state-of-the-art models including SAM (segment anything), T5, Code Whisper amongst others showing high numerical equivalency. Code: https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/MaxCode
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。