PyTorch-FX Shaper

Tools · 7

• 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