extract.py: shape parsing fails for PyTorch profiler tensor shape lists

[jay-tau]

Problem

profile.py writes shape_info as raw PyTorch profiler format (e.g. [[4096, 4096], [4096, 4096]] from str(evt.input_shapes)), but extract.py's parse_shape_info() only handles key=value format like M=4096, N=4096, K=4096.

This causes all profiled models to fall back to default shapes (e.g. 2048x2048x2048 for matmul), making extracted kernels not model-specific.

Affected Code

• extract.py parse_shape_info() (line 176) — regex only matches key=value pairs
• profile.py (line 633) — writes str(evt.input_shapes) which produces list-of-lists format

Reproduction

from extract import parse_shape_info
parse_shape_info("[[4096, 4096], [4096, 4096]]", "matmul")  # Returns None (should return {"M": 4096, "K": 4096, "N": 4096})

Proposed Fix

Add a fallback parser using ast.literal_eval that interprets PyTorch profiler tensor shape lists per op type:

• matmul: [[M,K],[K,N]] → {M, N, K} (also handles addmm [[N],[M,K],[K,N]])
• layernorm/softmax/rmsnorm/reduce: [[M,N],...] → {M, N}
• flash_attention/rotary_embedding: [[B,H,N,D],...] → {B, H, N, D}
• cross_entropy: [[batch,vocab],[batch]] → {batch, vocab}

This ensures extracted kernels use the actual model-specific shapes from profiling rather than hardcoded defaults.