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

• URL: http://arxiv.org/abs/2606.15994v1
• Date: Sun, 14 Jun 2026 19:41:57 GMT
• Title: Agentic Framework for Deep Learning workload migration via In-Context Learning
• Abstract summary: Large Language Models (LLMs) are prone to mistakes for exacting operations.<n>We propose a fully autonomous system that combines In-Context Learning (ICL) with oracle-driven self-ging.<n>Our pipeline achieves 91% numerical equivalence (compared to baseline: 9%, instruction + self-ging: 27%) on neural modules.
• Score: 0.5496900364175868
• License: http://creativecommons.org/licenses/by/4.0/
• 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 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

Related papers

• ARISE: A Repository-level Graph Representation and Toolset for Agentic Fault Localization and Program Repair [0.0]
  We present ARISE (Agentic Repository-level Issue Solving Engine), which augments an LLM-based agent with a multi-granularity program graph.<n>ARISE exposes this graph through a three-tier tool API, which brings data-flow slicing as a first-class, queryable agent primitive.<n>We evaluate on SWE-bench Lite (300 real GitHub issues, 11 Python repositories) using Qwen2.5-Coder-32B-Instruct as the backbone.
  arXiv  Detail & Related papers  (2026-05-04T19:59:23Z)
• VibeTensor: System Software for Deep Learning, Fully Generated by AI Agents [42.56489784841984]
  "fully generated" refers to code provenance: implementation changes were produced and applied as agent-proposed diffs.<n>We describe the architecture, summarize the workflow used to produce and validate the system, and evaluate the artifact.
  arXiv  Detail & Related papers  (2026-01-21T19:29:00Z)
• CEDAR: Context Engineering for Agentic Data Science [3.1662160826016756]
  CEDAR is an application for automating data science tasks with an agentic setup.<n>We show that these can be alleviated via effective context engineering.<n>Fault tolerance and context management are introduced via iterative code generation and smart history rendering.
  arXiv  Detail & Related papers  (2026-01-10T16:05:04Z)
• Semantic Soft Bootstrapping: Long Context Reasoning in LLMs without Reinforcement Learning [46.765013720309064]
  Long context reasoning in large language models (LLMs) has demonstrated enhancement of their cognitive capabilities via chain-of-thought (CoT) inference.<n>Training such models is usually done via reinforcement learning with verifiable rewards (RLVR) in reasoning based problems, like math and programming.<n>We propose textbfSemantic Soft Bootstrapping ( SSB), a self-distillation technique, in which the same base language model plays the role of both teacher and student, but receives different semantic contexts about the correctness of its outcome at training time.
  arXiv  Detail & Related papers  (2025-12-04T18:59:18Z)
• Agentic-R1: Distilled Dual-Strategy Reasoning [58.73951532294446]
  Current long chain-of-thought (long-CoT) models excel at mathematical reasoning but rely on slow and error-prone natural language traces.<n>We introduce a fine-tuning framework, DualDistill, that distills complementary reasoning strategies from multiple teachers into a unified student model.<n>Our method improves accuracy across a range of tasks, including both computation-intensive and standard benchmarks.
  arXiv  Detail & Related papers  (2025-07-08T06:35:16Z)
• SwingArena: Competitive Programming Arena for Long-context GitHub Issue Solving [90.32201622392137]
  We present SwingArena, a competitive evaluation framework for Large Language Models (LLMs)<n>Unlike traditional static benchmarks, SwingArena models the collaborative process of software by pairing LLMs as iterations, who generate patches, and reviewers, who create test cases and verify the patches through continuous integration (CI) pipelines.
  arXiv  Detail & Related papers  (2025-05-29T18:28:02Z)
• Generating Unseen Code Tests In Infinitum [1.0674604700001968]
  We present a method for creating benchmark variations that generalize across coding tasks and programming languages. We implement one benchmark, called textitauto-regression, for the task of text-to-code generation in Python.
  arXiv  Detail & Related papers  (2024-07-29T08:11:20Z)
• Get my drift? Catching LLM Task Drift with Activation Deltas [55.75645403965326]
  Task drift allows attackers to exfiltrate data or influence the LLM's output for other users.<n>We show that a simple linear classifier can detect drift with near-perfect ROC AUC on an out-of-distribution test set.<n>We observe that this approach generalizes surprisingly well to unseen task domains, such as prompt injections, jailbreaks, and malicious instructions.
  arXiv  Detail & Related papers  (2024-06-02T16:53:21Z)
• Allies: Prompting Large Language Model with Beam Search [107.38790111856761]
  In this work, we propose a novel method called ALLIES. Given an input query, ALLIES leverages LLMs to iteratively generate new queries related to the original query. By iteratively refining and expanding the scope of the original query, ALLIES captures and utilizes hidden knowledge that may not be directly through retrieval.
  arXiv  Detail & Related papers  (2023-05-24T06:16:44Z)
• Self-Prompting Large Language Models for Zero-Shot Open-Domain QA [67.08732962244301]
  Open-Domain Question Answering (ODQA) aims to answer questions without explicitly providing background documents. This task becomes notably challenging in a zero-shot setting where no data is available to train tailored retrieval-reader models. We propose a Self-Prompting framework to explicitly utilize the massive knowledge encoded in the parameters of Large Language Models.
  arXiv  Detail & Related papers  (2022-12-16T18:23:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.