PyTorch Optimization Agent
This challenge tasks you with building an intelligent agent focused on enhancing machine learning code performance. You will develop an agent capable of analyzing, refactoring, and generating PyTorch code for improved efficiency on various accelerators (e.g., TPUs, GPUs). The agent will leverage DeepSeek-R1 for its advanced code understanding and generation capabilities. LlamaIndex will be used to implement sophisticated RAG, providing context from PyTorch documentation, optimization guides, and existing codebases. Marvin will facilitate declarative AI functions and MCP-enabled tool integration, allowing the agent to interact with code linters, profilers, and even simulated CI/CD pipelines to validate and apply optimizations. The goal is to build a smart assistant for ML engineers, drastically reducing the manual effort in performance tuning.
What you are building
The core problem, expected build, and operating context for this challenge.
Develop a DeepSeek-R1 powered agent using LlamaIndex for RAG and Marvin for declarative MCP tool integration, focused on optimizing PyTorch code for AI chip efficiency.
How work is evaluated
The evaluation will assess the agent's capability to analyze, optimize, and generate PyTorch code effectively, focusing on performance gains, code quality, and the accuracy of explanations and tool interactions.
Challenge datasets
Public datasets, plus the private datasets you have access to.
What you should walk away with
Utilize DeepSeek-R1 for effective PyTorch code generation and refactoring.
Develop performance evaluation methods for optimized PyTorch code.
Implement advanced RAG pipelines for code optimization context.
Design a hybrid reasoning system for iterative code improvement.
Integrate ML engineering tools (profilers, linters) using Marvin's declarative AI and MCP.
Reference links and supporting material
Sample data for 2 tasks
Challenge setup
The evaluation will assess the agent's capability to analyze, optimize, and generate PyTorch code effectively, focusing on performance gains, code quality, and the accuracy of explanations and tool interactions.
Challenge input
{'pytorch_code': 'string', 'optimization_goal': 'lower inference latency'}
Agent execution
No agent tool is listed for this challenge.
Evaluated output
{'bottlenecks_identified': ['string'], 'proposed_optimizations': {'code_changes': 'string', 'explanation': 'string'}, 'tool_interactions_log': 'str...
- Generated or modified PyTorch code must be syntactically correct and runnable without errors.
- The agent's reasoning and code generation must demonstrate clear utilization of retrieved context from Llam...
- The agent must successfully invoke and interpret results from mock MCP-enabled tools (e.g., profiler, linte...
- Performance Improvement (Percentage) target: 20
- Python execution harness
Summary shows the first task, up to three tools, and up to three evidence items. Open the technical recipe to view all configured items.
View technical recipe
Configured tools
No tool records are attached.
Evaluation contract
- Generated or modified PyTorch code must be syntactically correct and runnable without errors.
- The agent's reasoning and code generation must demonstrate clear utilization of retrieved context from LlamaIndex RAG, with explanations referencing sources.
- The agent must successfully invoke and interpret results from mock MCP-enabled tools (e.g., profiler, linter) to inform its optimization or verification process.
Metrics
- Performance Improvement (Percentage) · Target 20
- Code Quality Score (Linter) · Target 90
- Explanation Coherence and Accuracy · Target 4
Workflow tasks
- Task 1
Input: {'pytorch_code': 'string', 'optimization_goal': 'lower inference latency'}
Output: {'bottlenecks_identified': ['string'], 'proposed_optimizations': {'code_changes': 'string', 'explanation': 'string'}, 'tool_interactions_log': 'string'}
Execution harness: Python
- Task 2
Input: {'task_description': 'string', 'performance_constraints': {'memory_limit_mb': 'int', 'target_flops': 'float'}, 'rag_context_keywords': ['attention mechanism', 'efficient pytorch']}
Output: {'generated_pytorch_module': 'string', 'design_justification': 'string', 'linter_report_status': 'string'}
Execution harness: Python
Recipe freeze marker
No recipe freeze time is recorded.