get_annotated_sourcesLABSCO SUMMARY
Shapes read off a real trace, then the rewrite checked.
Reading a transformer's forward pass usually means holding shapes in your head. This puts the measured shape on every line, and when your assistant rewrites matmuls and attention as einsum, the server checks the result against the original before you take it.
What it is
A shape-extraction and rewrite pipeline for PyTorch, exposed over MCP. It runs torch.export and ShapeProp across your nn.Module to capture the shape of every intermediate tensor, then returns your source annotated with those shapes and checks any einsum rewrite against the original.
What you get
- Per-module tensor shapes for every intermediate in the forward pass, not just the inputs
- Your source files back with those shapes inserted as inline comments on the lines they belong to
- Shape-annotated snippets for each class and function, which your assistant reads before it rewrites the code with einsum or einops
- Symbolic dimension names on request — pass a dim_names mapping and a shape reads (B, T) instead of (4, 16)
- A rewrite validated before you accept it — whole file, single module or single function
Requirements
The repository checked out locally with PyTorch, and a module that torch.export can trace — the shapes come from a real forward pass, not from static analysis. The server rewrites nothing itself: your assistant writes the einsum version and the server confirms it produces identical outputs. The library's separate comment and einsum modes call an LLM of their own and are not part of the MCP server.
Setup effort
Build from source — get the source, then follow the steps on this page and add the configuration for your client
USE CASES
01
Reading an unfamiliar model's forward pass
The shape of every intermediate tensor, captured from a real trace, so you stop inferring dimensions from variable names.
02
Annotating a file you have to maintain
Get the source back with shapes as inline comments on each line, so the notes stay with the code instead of in a notebook.
03
Converting attention maths to einsum safely
The rewrite is checked before you accept it, so collapsing reshapes into einsum does not quietly change the result.
04
Checking one function rather than a whole file
Validate a single rewritten function when only one block changed and a full-file pass is more than the edit deserves.
7 tools
Listed in project docs. Your client may show a different list.
get_fx_shapesvalidate_rewritevalidate_rewrite_functionvalidate_file_rewritesave_fixturesgenerate_test_filesBEFORE YOU INSTALL
What to know first
Tools you get
7 tools named in the project's documentation. The LLM rewriting modes and their OPENAI_API_KEY belong to the Python library, not the MCP server; the server extracts shapes and validates rewrites.
Requires
- Python 3.12 or later
- A PyTorch model to inspect, with a setup script that assigns model and example_args
Where it runs
The documented client runs the server on your machine.
INSTALLATION
Connect your client
Get the source, then follow the steps on this page and add the configuration for your client. This one is built from source rather than installed with a single command.
Server setup
TERMINAL
git clone https://github.com/Athe-kunal/Agent-Shaper
cd Agent-Shaper
python -m venv .venv
source .venv/bin/activate
pip install -e .Add the mcpServers block below to your Claude Code MCP config, pointing command at .venv/bin/python
MCP Clients
Claude Code
CLAUDE CODE · JSON
{
"mcpServers": {
"agent-shaper": {
"command": "<YOUR_PATH_TO_VENV>/bin/python",
"args": [
"-m",
"agent_shaper.mcp_server"
]
}
}
}
