Agentic Framework for Deep Learning workload migration via In-Context Learning

Abstract: 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 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

• The paper introduces an autonomous agentic pipeline that integrates in-context learning with execution oracle debugging for reliable PyTorch to JAX migration.
• It details a four-phase method, including ICL anchoring and hierarchical test harness generation, ensuring strict numerical equivalence (error < 1e-7).
• Empirical results show significant improvements, achieving 91-100% numerical equivalence and nearly complete repository migration validated by human experts.

Agentic Framework for Deep Learning Workload Migration via In-Context Learning

Motivation and Problem Definition

The migration of deep learning models from PyTorch's flexible, object-oriented paradigm to JAX’s functional, stateless ecosystem is essential for leveraging modern accelerators such as Google’s TPUs, yet remains a nontrivial, labor-intensive, and error-prone process. This migration is complicated by fundamental discrepancies: PyTorch's eager execution, dynamic tensor layouts (NCHW) and default fp32 precision contrast with JAX’s JIT compilation, NHWC layouts, bfloat16 precision, and compiler-intrinsic abstractions. Automated migration, via LLMs, is hampered by strict API desynchronization, frequent mathematical errors, and hallucinated code structure. The paper introduces an autonomous, agentic pipeline integrating ICL with oracle-driven dynamic self-debugging, targeting reliable cross-framework translation.

System Architecture

The proposed pipeline decomposes the migration task into four core phases: ICL anchoring, execution oracle generation, agentic translation with oracle-based test harnesses, and iterative self-debugging.

Figure 1: Four-phase pipeline: context anchoring, oracle extraction, core translation, and self-debugging to ensure agentic self-correction.

In-Context Learning Anchoring

ICL is operationalized by curating small, highly relevant reference sets comprising PyTorch source, idiomatic JAX translations, and test cases. Unlike retrieval-augmented methods, which introduce extraneous noise and aggravate model confusion ("Lost in the Middle" [liu2024lost]), ICL provides rich structural anchors allowing the LLM to compact demonstration vectors and enforce strict syntactic boundaries [hendel2023context, min2022rethinking]. This removes ambiguity in translation, curtails hallucination, and fosters task-aligned generation.

Execution Oracle Extraction

The system treats the original PyTorch implementation as the immutable mathematical ground truth. Profiling scripts pass random inputs through the source module, capturing initialized weights, intermediate activations, and output tensors, which are serialized into the "Oracle." This oracle artifact exposes hidden execution behaviors (silent padding, shape shifts), avoiding reliance on static analysis or LLM predictions, which are error-prone for high-dimensional tensor computation [dziri2023faith]. The oracle thus guarantees a fixed reference for downstream equivalence assertions.

Agentic Translation and Test Harness Generation

Leveraging ICL anchors and the dynamic oracle, the agent translates modules into flax.linen abstractions, then synthesizes test harnesses conditioned on oracle data. Verification proceeds hierarchically: compilation success, shape equivalence, and strict numerical equivalence (absolute error $< 10^{-7}$). Tests deserialize oracle .pkl files, mapping PyTorch weights to JAX, checking the model under realistic constraints. This method operationalizes the evaluation by execution paradigm [chen2021evaluating].

Iterative Self-Debugging

Failure modes in zero-shot and instruction-only setups include syntactic hallucinations, API signature mismatches, and state mapping errors. The agentic loop incorporates feedback from test harness executions—tracebacks, error logs, and detailed numerical differences—into the next agent input, guiding self-correction in structure and logic. Static instructions (PRNG key splitting constraints, layout normalization) support edge case mitigation, but empirical validation is achieved solely via iterative, oracle-conditioned debugging.

Empirical Results

The full pipeline achieves significant improvements over baselines and instruction/self-debugging-only methods. In controlled evaluations:

• Numerical equivalence on neural modules (Level 2): 91% (vs. baseline: 9%, instruction + self-debugging: 27%)
• Numerical equivalence on mathematical ops (Level 1): 100%
• Repository-level migration: Human experts validated near-complete task accomplishment (100% completeness for most repositories), with numerical equivalence reaching 100% for several, though below 85% for foundational, expansive Facebook Research models.
• Readability scores averaged 3.75–4.9, confirming practical usability.

Taxonomy of translation failures in baseline settings includes: failure to generate wrapper classes, misuse of PyTorch-specific kwargs, and mismatches in state dictionaries during weight transfer, underscoring the necessity of agentic, oracle-driven correction.

Implications and Future Directions

Practically, the agentic pipeline enables robust, scalable migration from PyTorch to JAX, unlocking optimization opportunities for accelerator-centric workloads and removing manual translation bottlenecks. Theoretical implications include advances in LLM grounding: by tightly integrating execution-based truth, the pipeline minimizes reward hacking and increases reliability in structurally complex translation tasks. The ablation studies dissect contributions: structural anchoring via ICL alone is insufficient for high-fidelity translation; empirical execution grounding is essential.

Unaddressed challenges include repository-level full automation, which demands elaborate dependency tracing and comprehensive module integration; additionally, the impact of varying ICL reference sets remains unexplored. Future research could extend agentic loops to other framework pairs and systematically optimize ICL context selection and oracle artifact construction. Integration with multi-agent collaboration models, as explored in related migration research [nikolov2026multi], presents further scalability prospects.

Conclusion

The agentic workflow, combining ICL-driven anchoring and execution oracle self-debugging, establishes a highly reliable and computationally tractable approach for translating deep learning workloads between PyTorch and JAX. Empirical data confirms substantial gains in numerical equivalence and translation completeness, validating the system’s design. The pipeline provides a generalizable blueprint for cross-framework migration, directly addressing critical engineering and research requirements in the context of AI accelerator optimization.