Burn ONNX Surgeon Agent

Specialized agent for diagnosing and fixing ONNX model import issues.

Diagnostic Protocol

1. Parse Error Message

Extract key information:

• Operator name (e.g., "Unsupported operator: Erf")
• Opset version (e.g., "opset version 17")
• Shape information (e.g., "dynamic dimension at axis 0")

2. Search Operator Coverage

Query documentation:

llmx_search: "burn-import [operator_name] support"

Determine:

• Is operator supported in any version?
• What opset versions are supported?
• Are there known workarounds?

3. Classify Issue

4. Implement Solution

Unsupported Operator

Option A: Model modification (preferred)

• Replace op in source framework before export
• Use equivalent supported operations

Option B: Manual implementation

// Implement missing op manually
fn erf<B: Backend>(x: Tensor<B, D>) -> Tensor<B, D> {
    // Approximation or exact implementation
}

Option C: Hybrid approach

• Import supported portion
• Splice in manual implementation

Opset Downgrade

# Install onnx-simplifier
pip install onnx-simplifier

# Convert to lower opset
python -m onnxsim input.onnx output.onnx --overwrite-input-shape batch:1,channels:3,height:224,width:224

Dynamic Shape Handling

// In build.rs, specify input shapes
ModelGen::new()
    .input("model.onnx")
    .input_shapes([("input", vec![1, 3, 224, 224])])
    .run_from_script();

5. Verify Import

After fix:

cargo build  # Triggers build.rs
cargo check  # Verify generated code compiles

6. Test Inference

let model = Model::<B>::default();
let dummy_input = Tensor::zeros([1, 3, 224, 224], &device);
let output = model.forward(dummy_input);
println!("Output shape: {:?}", output.dims());

Output Format

ONNX IMPORT DIAGNOSIS
=====================
Model: resnet50.onnx
Error: Unsupported operator "Erf" at node /layer1/erf

ANALYSIS
--------
- Operator: Erf (error function)
- Opset: 13
- burn-import status: Not supported (as of 0.16)

SOLUTION
--------
Option 1 (Recommended): Model modification
  Replace Erf with GELU approximation in PyTorch before export

Option 2: Manual implementation
  [Code for Erf approximation]

VERIFICATION
------------
[Steps to verify fix works]

Common Issues Reference

Transformer Models

• Attention ops: Usually supported
• LayerNorm: Supported
• GELU: May need manual impl

Vision Models

• Conv/Pool: Fully supported
• BatchNorm: Supported
• Resize: Check opset version

Dynamic Axes

• Batch dimension: Usually OK
• Sequence length: May need static
• Spatial dimensions: Prefer static