Function parse_shape_info

extract.py:176–207

Parse a shape_info string like "M=4096, N=4096, K=4096" into a dict. Handles various formats: - "M=4096, N=4096, K=4096" - "B=1, H=32, N=4096, D=128" - "batch=4096, vocab=32000" - "rows=4096, cols=4096" Returns None if parsing fails.

(shape_info_str: str, op_type: str)

# ---------------------------------------------------------------------------

def parse_shape_info(shape_info_str: str, op_type: str) -> Optional[Dict[str, int]]:
    """
    Parse a shape_info string like "M=4096, N=4096, K=4096" into a dict.

    Handles various formats:
      - "M=4096, N=4096, K=4096"
      - "B=1, H=32, N=4096, D=128"
      - "batch=4096, vocab=32000"
      - "rows=4096, cols=4096"

    Returns None if parsing fails.
    """
    if not shape_info_str or not isinstance(shape_info_str, str):
        return None

    # Match key=value pairs
    pairs = re.findall(r"([A-Za-z_][A-Za-z0-9_]*)\s*=\s*(\d+)", shape_info_str)
    if not pairs:
        return None

    raw = {k: int(v) for k, v in pairs}

    # Map to canonical bench.py keys using alias map
    alias_map = SHAPE_ALIAS_MAP.get(op_type, {})
    if alias_map:
        canonical = {}
        for k, v in raw.items():
            mapped_key = alias_map.get(k, k)
            canonical[mapped_key] = v
        return canonical
    else:
        return raw


def shape_to_display(shape: Dict[str, int]) -> str: