Trace is a Python library developed at Microsoft that implements end-to-end generative optimization for AI agents. The repository accompanies a research paper accepted to NeurIPS 2024 and provides a fully functional implementation of a platform for optimizing AI systems using general feedback such as numerical rewards, natural language text, or compiler errors. The library generalizes the backpropagation algorithm by capturing and propagating an AI system's execution trace, implemented as a PyTorch-like framework where users write Python code directly and use Trace primitives to optimize specific components.
The core functionality centers on two main primitives: node and bundle. The node primitive defines computation graph nodes with automatic tracing of subsequent operations, while the bundle primitive wraps Python functions to make them optimizable. Users mark components as trainable to indicate what the optimizer can modify, then apply familiar PyTorch-style syntax to conduct optimization. This approach enables flexible control over agent learning and allows seamless integration of optimization into existing Python workflows without requiring users to wrap code in strings or use specialized domain-specific languages.
Trace supports three distinct optimizers with different tradeoffs. OPRO leverages large language models as optimizers and offers speed advantages but lacks computation graph awareness. TextGrad performs automatic differentiation via text and supports computation graphs and library integration but operates at slower speeds. OptoPrime, Trace's proposed algorithm, utilizes the entire computational graph for parameter updates and achieves 2-3x faster performance than TextGrad while supporting code as executable functions and comprehensive library support. Users can seamlessly switch between optimizers depending on their specific requirements.
The repository demonstrates broad applicability across multiple domains through its tutorial collection. Beginner tutorials cover basic primitives and adaptive AI agents that learn from environment feedback, such as an LLM agent optimizing shot placement in Battleship. Intermediate examples showcase multi-agent collaboration in the VirtualHome environment and joint prompt and code optimization for NLP tasks on BigBench-Hard. Advanced tutorials demonstrate robotic arm control optimization through trajectory observation. The library requires Python 3.9 or higher and uses LiteLLM as the default LLM backend, with optional AutoGen support for backward compatibility.
GitGenius activity data reveals active development and maintenance patterns. The repository shows a median issue and pull request response latency of 2.1 hours with a mean of 70.4 hours across 23 tracked items, indicating responsive community engagement. The most active contributors include chinganc with 56 events, allenanie with 42 events, and doxav with 17 events. Bug reports represent the most common issue type with 5 tracked instances, followed by enhancement requests with 2 instances. The repository overlaps with related projects including microsoft/foundry-local, nousresearch/hermes-agent, and langchain-ai/langchain through shared contributors, indicating integration within a broader ecosystem of AI agent and optimization tools.
Recent developments highlight growing adoption and recognition. The Trace paper received featured status in the G-Research NeurIPS 2024 highlight review, and external research from Nvidia, Stanford, Visa, and Intel applied Trace to optimize mapper code for parallel programming, achieving 1.3x speedup in under 10 minutes compared to expert-optimized code. The project maintains active community channels including Discord and a mailing list for questions and bug reporting, with developers actively monitoring these communication channels.
