PyTorch-FX Shaper

none★ 1productivity

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

Install

How to run this server

[ { "cmd": "pip install -e", "imports": [] } ]

server path: pytorch-fx-shaper

Tools

What this server exposes

get_annotated_sources

Run a setup script, trace the model, return shape-annotated source for the requested files

get_fx_shapes

Return raw FX shape data for modules and functions

validate_rewrite

Check a rewritten nn.Module class produces identical outputs (atol=1e-5)

validate_rewrite_function

Check a rewritten standalone function produces identical outputs

validate_file_rewrite

Validate all classes in a rewritten file in one call

save_fixtures

Capture and persist forward-pass fixtures for later testing

generate_test_files

Generate pytest files that assert output identity against saved fixtures

Configuration

Environment & auth

authnone

envOPENAI_API_KEY

envOPENAI_BASE_URL

Resources