🤖 AI

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

Dit artikel stelt een autonoom agentisch framework voor dat gebruikmaakt van In-Context Learning en een executie-orakel afgeleid van werkelijke PyTorch-outputs om een uiterst betrouwbare, geautomatiseerde migratie van deep learning-modellen van PyTorch naar JAX te mogelijk maken, waarbij 91% numerieke equivalentie wordt bereikt door middel van iteratieve zelf-debugging.

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

Gepubliceerd 2026-06-16

📖 5 min leestijd🧠 Diepgaand

CC BY 4.0

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

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je een gezin probeert te verhuizen van een huis gebouwd van hout en spijkers (PyTorch) naar een huis dat volledig is gebouwd van glas en staal (JAX). Beide huizen hebben hetzelfde aantal kamers en dienen hetzelfde doel, maar de regels voor hoe ze gebouwd worden zijn totaal anders.

In het "hout"-huis kun je een spijker erin slaan en deze later weer veranderen; de structuur is flexibel. In het "glas"-huis moet alles vooraf gepland, rigide en onveranderlijk zijn. Als je probeert het glazen huis te bouwen met dezelfde timmertechnieken die je in het houten huis gebruikte, zal de hele constructie verbrijzelen.

Dit artikel beschrijft een nieuwe AI-bouwploeg (een "Agentic Framework") die ontworpen is om deze verhuizing te automatiseren. Hier is hoe ze het probleem hebben opgelost, met behulp van eenvoudige analogieën:

Het Probleem: De AI-Architect raakt in de war

Normaal gesproken, wanneer je een slimme AI (een Large Language Model) vraagt om code te vertalen van PyTorch naar JAX, gedraagt het zich als een architect die nog nooit een glazen huis heeft gezien. Het probeert de blauwdrukken te raden.

Het kan de wiskunde fout doen omdat het de complexe berekeningen niet in zijn hoofd kan "doen".
Het kan vergeten dat het glazen huis andere regels nodig heeft voor hoe kamers met elkaar verbonden zijn.
Het resultaat? Een gebouw dat er op papier goed uitziet, maar instort zodra je er probeert te wonen.

De Oplossing: Een Drieledig Veiligheidsnet

De auteurs hebben een systeem gebouwd dat de AI stopt met gokken en de AI dwingt om een strikt, geverifieerd proces te volgen. Ze noemen dit een "grounded generation" systeem.

1. De "Stijlgids" (In-Context Learning)

In plaats van de AI te laten raden hoe de code voor het glazen huis geschreven moet worden, gaf het team het een strikte stijlgids.

De Analogie: Stel je voor dat je de AI een fotoalbum geeft van de beste ooit gebouwde glazen huizen, waarin precies wordt getoond hoe de ramen en deuren geplaatst moeten worden.
Wat het doet: Deze "In-Context Learning" (ICL) fungeert als een anker. Het laat de AI zien: "Dit is hoe we code schrijven in JAX." Het voorkomt dat de AI hallucineert (dingen verzint) en houdt het op het juiste pad.

2. De "Waarheidsmachine" (De Oracle)

Dit is het belangrijkste onderdeel. De AI is slecht in het uitvoeren van complexe wiskunde in zijn hoofd. Daarom vroeg het team de AI niet om de getallen te berekenen. In plaats daarvan lieten ze de originele PyTorch-code eerst de wiskunde doen.

De Analogie: Voordat de AI probeert het glazen huis te bouien, laten ze het houten huis door een simulator lopen en maken ze een "snapshot" van exact hoe elke kamer zich gedraagt, hoe zwaar de meubels zijn en hoe het licht op de muren valt. Ze slaan deze snapshot op als een Oracle (een bron van absolute waarheid).
Wat het doet: De AI raadt de getallen niet. Het moet een glazen huis bouwen dat exact overeenkomt met deze snapshot. Als het glazen huis niet overeenkomt met de snapshot, is het fout.

3. De "Zelfcorrigerende Inspecteur" (Iteratieve Debugging)

De AI bouwt een concept van het glazen huis. Vervolgens controleert een robot-inspecteur het tegen de snapshot van de "Waarheidsmachine".

De Analogie: Als de inspecteur een raam vindt dat de verkeerde maat heeft of een deur die niet opengaat, zeggen ze niet alleen "Fout". Ze geven de blauwdruk terug aan de AI-architect en zeggen: "Kijk, het licht valt hier anders op de muur dan op de foto. Los het op."
Wat het doet: De AI probeert het opnieuw, leest de foutmelding en herstelt de code. Dit proces herhaalt zich totdat het glazen huis perfect overeenkomt met de snapshot.

De Resultaten: Van "Misschien" naar "Perfect"

Het team testte dit systeem op twee soorten taken:

Eenvoudige Wiskundige Operaties (Niveau 1): Zoals het verplaatsen van een enkele baksteen.
Complexe Neurale Modules (Niveau 2): Zoals het bouwen van een hele vleugel van het huis met complexe bedrading.

Zonder dit systeem (De Baseline):

De AI kreeg de wiskunde slechts 9% van de tijd goed voor complexe modules. Het was vooral aan het gokken.

Met alleen instructies (Zonder Oracle):

De AI werd beter in het volgen van regels, maar maakte de wiskunde nog steeds 73% van de tijd fout. Het bouwde een huis dat er goed uitzag, maar niet functioneerde.

Met het Volledige Systeem (De Methode van het Papier):

Eenvoudige taken: 100% succes.
Complexe taken: 91% succes.

Het systeem migreerde beroemde modellen zoals SAM (die computers helpt "beelden te zien"), T5 (een tekstvertaler) en Code Whisper (een assistent bij het programmeren) met hoge nauwkeurigheid van PyTorch naar JAX.

De Kernboodschap

Het paper bewijst dat je niet simpelweg een AI kunt vragen om complexe code te "vertalen" en dan maar hopen op het beste. Je moet het voorzien van:

Voorbeelden van hoe het juist moet (De Stijlgids).
Bewijs van hoe het resultaat eruit moet zien (De Oracle).
Een kans om fouten te herstellen op basis van dat bewijs (De Inspecteur).

Door deze drie te combineren, creëerden ze een betrouwbare, geautomatiseerde manier om deep learning-modellen tussen verschillende frameworks te verplaatsen zonder dat een mens elke regel code hoeft te controleren.

Technische Samenvatting: Agentic Framework voor de Migratie van Deep Learning Workloads via In-Context Learning

1. Probleemstelling

De migratie van deep learning-modellen van het PyTorch-paradigma (eager, objectgeoriënteerd en mutabele staat) naar het JAX-paradigma (functioneel, staatloos en JIT-gecompileerd/XLA) vormt een aanzienlijke engineering-bottleneck. Deze transitie is cruciaal voor het benutten van hardware-specifieke optimalisaties, zoals Google's Tensor Processing Units (TPUs), die superieure kostenefficiëntie en snelheden voor matrixvermenigvuldiging bieden. De handmatige migratie is echter foutgevoelig vanwege fundamentele verschillen in:

Executiemodellen: PyTorch's eager execution versus JAX's Just-In-Time (JIT) compilatie.
State Management: PyTorch's mutabele objectstaten versus JAX's functionele, staatloze flax.linen abstracties.
Control Flow: Standaard Python-loops versus compiler-vriendelijke structuren zoals jax.lax.scan.
Numerieke Discrepanties: Verschillen in standaard tensor-layouts (NCHW versus NHWC), precisies (fp32 versus bfloat16) en padding-gedrag in attention-lagen.

Hoewel Large Language Models (LLMs) bekwaamheid hebben getoond in algemene codegeneratie, worstelen zij met strikte framework-translatie. Zonder betrouwbare verificatie hebben LLM's de neiging tot hallucinaties wat betreft wiskundige operaties en semantische logica, met name met betrekking tot hoogdimensionale tensor-arithmetiek, wat leidt tot code die weliswaar compileert, maar niet numeriek equivalente resultaten produceert.

2. Methodologie

De auteurs stellen een volledig autonoom, agentic pijplijn voor die ontworpen is om PyTorch-modules naar JAX te vertalen. Het systeem behandelt framework-translatie als een "grounded generation problem" en maakt gebruik van een vierfasen workflow (geïllustreerd in Figuur 1 van het paper) die In-Context Learning (ICL) combineert met een execution-based oracle.

Fase 1: Verankering via In-Context Learning (ICL)

Om het "Lost in the Middle"-fenomeen te mitigeren en ruis door brede retrieval-methoden (zoals RAG) te verminderen, hanteert het systeem een strikt gecontroleerde ICL-strategie. In plaats van grote, complexe codefragmenten op te halen, wordt de context beperkt tot een kleine, dichte set referenties die het volgende bevat:

De bron PyTorch-module.
De idiomatische JAX-translatie.
Relevante testgevallen.
Deze voorbeelden dienen als structurele ankers, die de LLM begeleiden bij het adopteren van de juiste syntactische formattering en functionele patronen zonder het contextvenster te overbelasten.

Fase - 2: De Dynamische Execution Oracle

In het besef dat LLM's niet betrouwbaar hoogdimensionale tensor-arithmetiek kunnen afleiden, stelt het systeem een onveranderlijke execution oracle vast.

De bron PyTorch-module wordt uitgevoerd met willekeurige input-tensors.
De resulterende dynamische computationele staat—inclusief geïnitialiseerde gewichten (state_dict), input-tensors en output-activaties—wordt geserialiseerd naar een vast artefact (.pkl bestanden).
Deze oracle legt gedragingen vast die statische analyse mist, zoals stille padding-methoden en dynamische vormveranderingen, en biedt daarmee de absolute wiskundige grondwaarheid voor de doel-JAX-code.

Fase 3: Agentic Translatie en Testgeneratie

De agent vertaalt de PyTorch-bron naar flax.linen abstracties, begrensd door de ICL-ankers. Cruciaal is dat de agent de taak krijgt om custom test-harnesses te genereren op basis direct van de Fase 2 Oracle. De evaluatie volgt een hiërarchisch driestaps proces:

Compilation Check: Verifieert of de JAX-module initialiseert en traceert zonder runtime-fouten.
Shape Correctness: Bevestigt dimensionale consistentie en de afhandeling van layout-verschuivingen (bijv. NCHW naar NHWC).
Numerieke Equivalentie: Het script deserialiseert de Oracle-gewichten, mapt deze naar het JAX-model en controleert op numerieke equivalentie met behulp van numpy.allclose met een strikte absolute foutenlimiet van $10^{-7}$ . Deze worden "Silver" testgevallen genoemd.

Fase 4: Iteratieve Self-Debugging

Indien de gegenereerde code faalt voor de Oracle-geconditioneerde tests, gaat het systeem een zelfcorrectie-loop binnen. Executie-feedback, inclusief stack traces, standaardfouten en precieze numerieke verschillen, wordt teruggevoerd in de context van de agent. Geleid door de ICL-templates en deze objectieve feedback, verfijert de LLM iteratief de code totdat deze voldoet aan de strikte wiskundige restricties.

3. Belangrijkste Bijdragen

Het paper schetst drie primaire bijdragen:

Structurele Verankering via ICL: Het aantonen dat een kleine, zorgvuldig gecureerde set referentie-translaties en tests het generatieproces effectief verankert, wat het risico op hallucinaties aanzienlijk verlaagt vergeleken met brede retrieval of instructie-only benaderingen.
De Dynamische Execution Oracle: Het introduceren van een methode om de rekenkundige beperkingen van LLM's te omzeilen door de bron PyTorch-implementatie als de grondwaarheid te gebruiken. Dit sluit aan bij de verschuiving naar execution-based code-evaluatie, waarbij wiskundige getrouwheid wordt gewaarborgd door echte berekening in plaats van theoretische schatting.
Ablatie-begeleiding voor Agentic Workflows: Het leveren van empirisch bewijs via ablatie-studies die de impact van instructies, ICL en execution-conditioned self-debugging isoleren, wat een blauwdruk biedt voor toekomstige agentic code-translatie systemen.

4. Resultaten

Het systeem werd geëvalueerd op een holdout-set gecategoriseerd in drie moeilijkheidsgraden:

Niveau 1: Fundamentele wiskundige operaties (9 taken).
Niveau 2: Neurale modules en architecturale abstracties (11 taken, waaronder MLPs, Transformers, GANs, etc.).
Niveau 3: Volledige repository-migratie (10 repositories, waaronder SAM, T5, Code Whisper, FinGPT).

Kwantitatieve Prestaties (Tabel 1):

Baseline (Alleen Instructie): Bereikte slechts 9% numerieke equivalentie op Niveau 2 modules.
Instructie + Self-Debugging: Verbeterde naar 27% numerieke equivalentie op Niveau 2, maar shape-matching bleef steken op 55%.
Volledige Pijplijn (ICL + Oracle + Self-Debugging): Bereikte 100% numerieke equivalentie op Niveau 1 en 91% op Niveau 2.

Repository Migratie (Tabel 2):
Over 10 Niveau 3 repositories bereikte de volledige pijplijn 100% volledigheid in code-conversie. Numerieke equivalentie varieerde per modelcomplexiteit:

Hoge Getrouwheid (93%–100%): Code Whisper, Multimodal Transformer, Time-series forecasting, StyleGAN2, T5, Two-tower modellen.
Gemiddelde Getrouwheid: FinGPT (86%), SAM2 (80%), DETR (57%).
Leesbaarheid: Menselijke evaluatoren beoordeelden de leesbaarheid van de geconverteerde bestanden tussen 3.75 en 4.9 op een 5-punts Likert-schaal.

Falanalyse:
Het paper identificeert drie hoofdfoutmodi in baselines: Syntactische Hallucinaties (ontbrekende wrapper classes), API Signature Mismatches (onjuiste kwargs) en State Mapping Discrepanties (KeyErrors tijdens weight transfer). De iteratieve debugging-loop van de volledige pijplijn adresseert deze specif gezien door de agent te dwingen de structurele en weight-mapping logica af te stemmen op de Oracle.

5. Betekenis en Claims

De auteurs beweren dat hun gegronde, agent-gestuurde pijplijn een zeer betrouwbaar en schaalbaar blauwdruk biedt voor cross-framework migratie. De betekenis van het werk ligt in:

Elimineren van Handmatige Testen: Het systeem elimineert de noodzaak voor handmatige testcreatie door "Silver" testgevallen te genereren die gegrond zijn in execution oracles.
Overwinnen van LLM-beperkingen: Door de rol van de LLM te ontkoppelen van wiskundige deductie (die wordt afgehandeld door de Oracle) en te focussen op structurele translatie en debugging, bereikt het systeem hoge-getrouwheid resultaten waar pure LLM-benaderingen falen.
Schaalbaarheid: De aanpak slaagde erin complexe, state-of-the-art modellen (bijv. SAM, T5) te migreren met een hoge numerieke equivalentie, wat het nut valideert voor real-world engineering uitdagingen.

Beperkingen:
De auteurs merken bescheiden twee huidige beperkingen op:

De methode is nog niet uitgebreid naar volledig automatische repository-niveau migraties die volledige structurele afhankelijkheidsbeheer vereisen.
De specifieke impact van variërende voorbeelden gebruikt in de ICL-context is nog niet systematisch onderzocht.

Ondanks deze beperkingen concludeert het paper dat het combineren van structurele verankering met dynamische, execution-based feedback effectief de bottleneck van framework-translatie aanpakt, waardoor de migratie van deep learning workloads naar geoptimaliseerde JAX-omgevingen met minimale handmatige interventie mogelijk wordt.