🤖 AI

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

This paper proposes an autonomous agentic framework that leverages In-Context Learning and an execution oracle derived from actual PyTorch outputs to enable highly reliable, automated migration of deep learning models from PyTorch to JAX, achieving 91% numerical equivalence through iterative self-debugging.

Original authors: Qiyue Liang, Steven Ingram, George Vanica, Andi Gavrilescu, Newfel Harrat, Hassan Sipra, Sethuraman Sankaran

Published 2026-06-16

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Qiyue Liang, Steven Ingram, George Vanica, Andi Gavrilescu, Newfel Harrat, Hassan Sipra, Sethuraman Sankaran

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to move a family from a house built with wood and nails (PyTorch) into a house built entirely of glass and steel (JAX). Both houses have the same number of rooms and serve the same purpose, but the rules for how they are built are completely different.

In the "wood" house, you can hammer a nail in, change it later, and the structure is flexible. In the "glass" house, everything must be pre-planned, rigid, and immutable. If you try to build the glass house using the same hammering techniques you used for the wood, the whole thing will shatter.

This paper describes a new AI construction crew (an "Agentic Framework") designed to automate this move. Here is how they solved the problem, using simple analogies:

The Problem: The AI Architect Gets Confused

Usually, when you ask a smart AI (a Large Language Model) to translate code from PyTorch to JAX, it acts like an architect who has never seen a glass house. It tries to guess the blueprints.

It might get the math wrong because it can't "do" the complex calculations in its head.
It might forget that the glass house needs different rules for how rooms connect.
The result? A building that looks okay on paper but collapses when you try to live in it.

The Solution: A Three-Part Safety Net

The authors built a system that stops the AI from guessing and forces it to follow a strict, verified process. They call this a "grounded generation" system.

1. The "Style Guide" (In-Context Learning)

Instead of letting the AI guess how to write the glass house code, the team gave it a strict style guide.

The Analogy: Imagine giving the AI a photo album of the best glass houses ever built, showing exactly how the windows and doors should be placed.
What it does: This "In-Context Learning" (ICL) acts as an anchor. It shows the AI, "This is how we write code in JAX." It prevents the AI from hallucinating (making things up) and keeps it on the right track.

2. The "Truth Machine" (The Oracle)

This is the most important part. The AI is bad at doing complex math in its head. So, the team didn't ask the AI to calculate the numbers. Instead, they let the original PyTorch code do the math first.

The Analogy: Before the AI tries to build the glass house, they run the wood house through a simulator and take a "snapshot" of exactly how every room behaves, how heavy the furniture is, and how the light hits the walls. They save this snapshot as an Oracle (a source of absolute truth).
What it does: The AI doesn't guess the numbers. It has to build a glass house that matches this snapshot exactly. If the glass house doesn't match the snapshot, it's wrong.

3. The "Self-Correcting Inspector" (Iterative Debugging)

The AI builds a draft of the glass house. Then, a robot inspector checks it against the "Truth Machine" snapshot.

The Analogy: If the inspector finds a window that is the wrong size or a door that doesn't open, they don't just say "Fail." They hand the blueprint back to the AI architect and say, "Look, the light hits this wall differently than in the photo. Fix it."
What it does: The AI tries again, reads the error, and fixes the code. It loops this process until the glass house matches the snapshot perfectly.

The Results: From "Maybe" to "Perfect"

The team tested this system on two types of tasks:

Simple Math Operations (Level 1): Like moving a single brick.
Complex Neural Modules (Level 2): Like building a whole wing of the house with complex wiring.

Without this system (The Baseline):

The AI got the math right only 9% of the time for complex modules. It was mostly guessing.

With just instructions (No Oracle):

The AI got better at following rules but still got the math wrong 73% of the time. It was building a house that looked right but didn't work.

With the Full System (The Paper's Method):

Simple tasks: 100% success.
Complex tasks: 91% success.

The system successfully migrated famous models like SAM (which helps computers "see" images), T5 (a text translator), and Code Whisper (a coding assistant) from PyTorch to JAX with high accuracy.

The Bottom Line

The paper proves that you can't just ask an AI to "translate" complex code and hope for the best. You have to give it:

Examples of how to do it right (The Style Guide).
Proof of what the result should look like (The Oracle).
A chance to fix its mistakes based on that proof (The Inspector).

By combining these three, they created a reliable, automated way to move deep learning models between different frameworks without needing a human to check every single line of code.

Technical Summary: Agentic Framework for Deep Learning Workload Migration via In-Context Learning

1. Problem Statement

The migration of deep learning models from PyTorch's eager, object-oriented, and mutable state paradigm to JAX's functional, stateless, and JIT-compiled (XLA) paradigm is a significant engineering bottleneck. This transition is critical for leveraging hardware-specific optimizations, such as Google's Tensor Processing Units (TPUs), which offer superior cost-efficiency and matrix multiplication speeds. However, manual migration is error-prone due to fundamental differences in:

Execution Models: PyTorch's eager execution vs. JAX's Just-In-Time (JIT) compilation.
State Management: PyTorch's mutable object states vs. JAX's functional, stateless flax.linen abstractions.
Control Flow: Standard Python loops vs. compiler-friendly structures like jax.lax.scan.
Numerical Discrepancies: Differences in default tensor layouts (NCHW vs. NHWC), precisions (fp32 vs. bfloat16), and padding behaviors in attention layers.

While Large Language Models (LLMs) have shown proficiency in general code generation, they struggle with strict framework translation. Without reliable verification, LLMs tend to hallucinate mathematical operations and semantic logic, particularly regarding high-dimensional tensor arithmetic, leading to code that compiles but fails to produce numerically equivalent results.

2. Methodology

The authors propose a fully autonomous, agentic pipeline designed to translate PyTorch modules to JAX. The system treats framework translation as a "grounded generation problem," utilizing a four-phase workflow (illustrated in Figure 1 of the paper) that combines In-Context Learning (ICL) with an execution-based oracle.

Phase 1: Anchoring via In-Context Learning (ICL)

To mitigate the "Lost in the Middle" phenomenon and reduce noise from broad retrieval methods (like RAG), the system employs a tightly controlled ICL strategy. Instead of retrieving large, complex code fragments, the context is limited to a small, dense set of references containing:

The source PyTorch module.
The idiomatic JAX translation.
Relevant test cases.
These examples serve as structural anchors, guiding the LLM to adopt correct syntactic formatting and functional patterns without overwhelming the context window.

Phase 2: The Dynamic Execution Oracle

Recognizing that LLMs cannot reliably deduce high-dimensional tensor arithmetic, the system establishes an unchangeable execution oracle.

The source PyTorch module is executed with random input tensors.
The resulting dynamic computational state—including initialized weights (state_dict), input tensors, and output activations—is serialized into a fixed artifact (.pkl files).
This oracle captures behaviors that static analysis misses, such as silent padding methods and dynamic shape changes, providing the absolute mathematical ground truth for the target JAX code.

Phase 3: Agentic Translation and Test Generation

The agent translates the PyTorch source into flax.linen abstractions, constrained by the ICL anchors. Crucially, the agent is tasked with generating custom test harnesses based directly on the Phase 2 Oracle. The evaluation follows a hierarchical three-stage process:

Compilation Check: Verifies the JAX module initializes and traces without runtime errors.
Shape Correctness: Confirms dimensional consistency, handling layout shifts (e.g., NCHW to NHWC).
Numerical Equivalence: The script deserializes the Oracle weights, maps them to the JAX model, and checks for numerical equivalence using numpy.allclose with a strict absolute error limit of $10^{-7}$ . These are termed "Silver" test cases.

Phase 4: Iterative Self-Debugging

If the generated code fails the Oracle-conditioned tests, the system enters a self-correction loop. Execution feedback, including stack traces, standard errors, and precise numerical differences, is fed back into the agent's context. Guided by the ICL templates and this objective feedback, the LLM iteratively refines the code until it passes the strict mathematical constraints.

3. Key Contributions

The paper outlines three primary contributions:

Structural Anchoring via ICL: Demonstrating that a small, carefully curated set of reference translations and tests effectively anchors the generation process, significantly lowering the risk of hallucination compared to broad retrieval or instruction-only approaches.
The Dynamic Execution Oracle: Introducing a method to bypass LLM arithmetic limitations by using the source PyTorch implementation as the ground truth. This aligns with the shift toward execution-based code evaluation, ensuring mathematical fidelity through real computation rather than theoretical estimation.
Ablation Guidance for Agentic Workflows: Providing empirical evidence through ablation studies that isolates the impact of instructions, ICL, and execution-conditioned self-debugging, offering a blueprint for future agentic code-translation systems.

4. Results

The system was evaluated on a holdout set categorized into three difficulty levels:

Level 1: Fundamental mathematical operations (9 tasks).
Level 2: Neural modules and architectural abstractions (11 tasks, including MLPs, Transformers, GANs, etc.).
Level 3: Full repository migration (10 repositories, including SAM, T5, Code Whisper, FinGPT).

Quantitative Performance (Table 1):

Baseline (Instruction Only): Achieved only 9% numerical equivalence on Level 2 modules.
Instruction + Self-Debugging: Improved to 27% numerical equivalence on Level 2, but shape matching remained at 55%.
Full Pipeline (ICL + Oracle + Self-Debugging): Achieved 100% numerical equivalence on Level 1 and 91% on Level 2.

Repository Migration (Table 2):
Across 10 Level 3 repositories, the full pipeline achieved 100% completeness in code conversion. Numerical equivalence varied by model complexity:

High Fidelity (93%–100%): Code Whisper, Multimodal Transformer, Time-series forecasting, StyleGAN2, T5, Two-tower models.
Moderate Fidelity: FinGPT (86%), SAM2 (80%), DETR (57%).
Readability: Human evaluators rated the readability of converted files between 3.75 and 4.9 on a 5-point Likert scale.

Failure Analysis:
The paper identifies three main failure modes in baselines: Syntactic Hallucinations (missing wrapper classes), API Signature Mismatches (incorrect kwargs), and State Mapping Discrepancies (KeyErrors during weight transfer). The full pipeline's iterative debugging loop specifically addresses these by forcing the agent to resolve structural and weight-mapping logic against the Oracle.

5. Significance and Claims

The authors claim that their grounded, agent-driven pipeline offers a highly reliable and scalable blueprint for cross-framework migration. The significance of the work lies in:

Eliminating Manual Testing: The system removes the need for manual test creation by generating "Silver" test cases grounded in execution oracles.
Overcoming LLM Limitations: By decoupling the LLM's role from mathematical deduction (handled by the Oracle) and focusing it on structural translation and debugging, the system achieves high-fidelity results where pure LLM approaches fail.
Scalability: The approach successfully migrated complex, state-of-the-art models (e.g., SAM, T5) with high numerical equivalence, validating its utility for real-world engineering challenges.

Limitations:
The authors modestly note two current limitations:

The method has not yet been expanded to fully automatic repository-level migrations requiring complete structural dependency management.
The specific impact of varying the examples used in the ICL context has not been systematically explored.

Despite these limitations, the paper concludes that combining structural anchoring with dynamic, execution-based feedback effectively addresses the framework translation bottleneck, enabling the migration of deep learning workloads to optimized JAX environments with minimal manual intervention.