PyTorch-FX Shaper

The MCP server provides shape of tensors to convert PyTorch code to einsum and einops

Agent Shaper extracts per-module tensor shape metadata from any PyTorchnn.Moduleand uses it to annotate source files — either with descriptive shape comments or by rewriting operations astorch.einsum/einops.
- Shape extraction(fx_utils/get_fx_data.py) — runstorch.export+ShapePropon your module to capture the shape of every intermediate tensor in every workspace-defined module's forward pass, not just the inputs.
- Manual annotation(fx_utils/manual_annotate.py) — inserts the extracted shapes as inline comments at the end of each relevant source line.
- LLM annotation(fx_utils/llm_annotate.py) — feeds each manually-annotated module class to an LLM in parallel. Two modes:

- COMMENT— rewrites comments with descriptive dimension names (e.g.batch_size,seq_len,n_embd) and plain-English explanations of each transformation.
- EINSUM— rewrites the entire module replacing matmuls and attention operations withtorch.einsum, collapsing intermediate reshapes where possible.

All modules in a file are processed in parallel viaasyncio. Everything outside the module classes (imports, dataclasses, config objects) is preserved unchanged in the output file.

agent_shaper/ fx_utils/ get_fx_data.py # shape extraction via torch.export + ShapeProp manual_annotate.py # inline shape comment insertion llm_annotate.py # LLM-powered rewrite (COMMENT or EINSUM mode) diff_viewer.py # Streamlit diff UI mcp_server.py # MCP tool server for AI-assisted rewrites examples/ transformer/ # GPT-2, LLaMA, Qwen3, Swin Transformer model.py # nanoGPT reference implementation model_einsum.py # einsum-rewritten version llama.py / llama_einsum.py qwen3.py / qwen3_einsum.py swin_transformer.py / swin_transformer_einsum.py run_transformer.py # example forward pass alignment/ # standalone alignment loss references alignment_losses.py # DPO/IPO/SimPO losses (reference einsum style) dpo_losses_einsum.py jsd.py / jsd_einsum.py opd.py / opd_einsum.py direct_alignment/ # full training-ready RLHF loss implementations loss.py # DPO, IPO, SimPO, ORPO, KTO, APO-zero, APO-down train.py # training loop data.py # preference dataset loading config.py # training configuration

Each_einsum.pyfile is the einsum/einops-rewritten counterpart of the original, validated to produce numerically identical outputs (atol=1e-5).

python -m venv .venv source .venv/bin/activate pip install -e .

{ "mcpServers": { "agent-shaper": { "command": "/path/to/.venv/bin/python", "args": ["-m", "agent_shaper.mcp_server"] } } }

- model: annn.Moduleinstance to trace
- example_args: a tuple of example tensors matchingforward()'s signature
- dim_names(optional):dictmapping symbolic dim names to integer values (e.g.{"B": 4, "T": 16}) so shapes show(B, T)instead of(4, 16)

import torch from agent_shaper.transformer.model import GPT, GPTConfig from agent_shaper.fx_utils.manual_annotate import annotate_module_source cfg = GPTConfig(block_size=32, vocab_size=256, n_layer=2, n_head=2, n_embd=64, dropout=0.0, bias=True) B, T = 2, 16 example_args = ( torch.randint(0, cfg.vocab_size, (B, T), dtype=torch.long), torch.randint(0, cfg.vocab_size, (B, T), dtype=torch.long), ) annotated = annotate_module_source( GPT(cfg), example_args, dim_names={"B": B, "T": T}, output_dir="annotated_output", # writes annotated files here; omit to just get the dict back )

python -m agent_shaper.fx_utils.manual_annotate

export OPENAI_API_KEY=sk-... export OPENAI_MODEL=gpt-4o export OPENAI_BASE_URL=https://your-proxy/v1 # optional; omit for default OpenAI

import asyncio, torch from agent_shaper.transformer.model import GPT, GPTConfig from agent_shaper.fx_utils.llm_annotate import llm_annotate_module_source, AnnotationMode cfg = GPTConfig(block_size=32, vocab_size=256, n_layer=2, n_head=2, n_embd=64, dropout=0.0, bias=True) B, T = 2, 16 example_args = ( torch.randint(0, cfg.vocab_size, (B, T), dtype=torch.long), torch.randint(0, cfg.vocab_size, (B, T), dtype=torch.long), ) async def main(): annotated = await llm_annotate_module_source( GPT(cfg), example_args, mode=AnnotationMode.COMMENT, # or AnnotationMode.EINSUM dim_names={"B": B, "T": T}, output_dir="llm_annotated_output", ) asyncio.run(main())

python -m agent_shaper.fx_utils.llm_annotate

.venv/bin/streamlit run agent_shaper/fx_utils/diff_viewer.py

dim_names = {"B": 2, "T": 16, "n_embd": 64, "n_head": 2}

from agent_shaper.fx_utils import get_module_shapes, TensorInfo module_infos = get_module_shapes(model, example_args, dim_names=dim_names) for info in module_infos: print(info.class_name, info.source_file, info.line_start, info.line_end) for t in info.tensors: print(" ", t.name, t.shape, t.annotated_shape)

- examples/transformer/— GPT-2, LLaMA, Qwen3, Swin Transformer. Each model has a_einsum.pycounterpart rewritten withtorch.einsumandeinops.
- examples/alignment/— standalone alignment loss functions (DPO, IPO, SimPO, JSD, OPD) in both original and einsum form.
- examples/direct_alignment/— production-style RLHF loss implementations (DPO, cDPO, IPO, SimPO, ORPO, KTO, APO-zero, APO-down) asnn.Moduleclasses with a full training loop. Theloss.pyfile useseinops.einsumandeinops.reducethroughout, validated against the gather-based originals at atol=1e-5.

All*_einsum.pyrewrites were validated using thevalidate_rewrite/validate_rewrite_functionMCP tools.

Solve optimization problems formulated by an LLM using the on-device Gurobi solver.

Train a Linear Regression model by uploading a CSV dataset file, demonstrating an end-to-end machine learning workflow.

Integrates with MLflow, enabling AI assistants to interact with experiments, runs, and registered models.

Neo is the first autonomous AI engineering agent that plans, researches and executes multi-step work for tasks such as building AI agents, AI model fine-tuning and evals, and ML pipelines; using your codebase, data, and experiments to ship faster with inspectable artifacts. It can reason over your repository, wire tools and retrieval, debug training runs, and help you develop production-ready AI workflows

A reasoning-first MCP middleware that uses heuristics and Neural BERT models to distil context and eliminate noise

An MCP server for accessing tidymodels GitHub information and generating code.

//beforeyouship - LLM cost modeling from your editor

Model the realistic monthly cost of an LLM app before you build it. Not a token calculator: retries, prompt caching, batch discounts, infra overhead, and 3×/10× growth are modeled in, across GPT-5.x, Claude, Gemini, DeepSeek, and more. Works without a key. Connect and ask — demo mode covers the six free-tier models. A Pro API key (beforeyouship.dev) unlocks the full 18-model catalog.

Recognize and extract text from handwritten documents using the Handwriting OCR service.

Perform advanced CSV analysis and generate insights using Google's Gemini AI. Requires Gemini and Plotly API keys.

MCP server that gives AI assistants on-demand access to 1,500+ amCharts docs, ~300 code examples, and 1000+ class API references.