fix_octen_dml.py

"""
Fix Octen-Embedding-0.6B ONNX for DirectML compatibility.

Root cause: dynamo torch.onnx.export generates `val_41 = [-1]` used in Reshape
shapes for multi-head attention Q/K/V projections. ONNX resolves -1 differently
per input (16 for Q with 2048-dim output, 8 for K/V with 1024-dim output).
DirectML's execution provider needs concrete shape values at graph-capture time
and crashes on symbolic -1 dims.

Fix: Replace val_41=[-1] with three concrete head-count constants:
  - val_41_q = [16]  (Q projection: num_attention_heads=16)
  - val_41_k = [8]   (K projection: num_key_value_heads=8)
  - val_41_v = [8]   (V projection: num_key_value_heads=8)

Create separate Concat nodes for each shape and reconnect the Reshape consumers
based on the Q/K/V naming pattern:
  - node_view{4L+0} (Q)  -> val_50_q
  - node_view{4L+1} (K)  -> val_50_k
  - node_view{4L+2} (V)  -> val_50_v
"""

import sys
import os
import shutil
from pathlib import Path

import onnx
import numpy as np
from onnx import helper, TensorProto

# octen-embedding-0.6b attention heads
NUM_Q_HEADS = 16
NUM_KV_HEADS = 8
HEAD_DIM = 128


def patch_octen_for_directml(input_path: str, output_path: str):
    print(f"Loading: {input_path}")
    m = onnx.load(input_path)
    g = m.graph
    initials = {i.name: i for i in g.initializer}

    # --- 1. Create new head-count initializers ---
    def make_int64_initializer(name: str, value: int) -> TensorProto:
        return helper.make_tensor(name, TensorProto.INT64, [1], [value])

    new_inits = [
        make_int64_initializer("val_41_q", NUM_Q_HEADS),
        make_int64_initializer("val_41_k", NUM_KV_HEADS),
        make_int64_initializer("val_41_v", NUM_KV_HEADS),
    ]

    # Add them to the graph (before the original val_41 so they're available)
    val_41_idx = next(i for i, init in enumerate(g.initializer) if init.name == "val_41")
    for offset, init in enumerate(new_inits):
        g.initializer.insert(val_41_idx + offset, init)

    # --- 2. Create new Concat nodes for val_50_q, val_50_k, val_50_v ---
    # Original: val_50 = Concat([val_0 (batch), val_1 (seq), val_41 (-1), val_49 (128)])
    # New:      val_50_q = Concat([val_0, val_1, val_41_q(16), val_49])
    #           val_50_k = Concat([val_0, val_1, val_41_k(8),  val_49])
    #           val_50_v = Concat([val_0, val_1, val_41_v(8),  val_49])
    #
    # The original val_50 Concat has attribute axis=0
    val_50_node = next(n for n in g.node if "val_50" in n.output)

    for suffix, val_name in [("q", "val_41_q"), ("k", "val_41_k"), ("v", "val_41_v")]:
        new_concat = helper.make_node(
            "Concat",
            inputs=[val_50_node.input[0], val_50_node.input[1], val_name, val_50_node.input[3]],
            outputs=[f"val_50_{suffix}"],
            name=f"node_val_50_{suffix}",
            axis=0,
        )
        # Insert after the original val_50 Concat
        val_50_pos = next(i for i, n in enumerate(g.node) if n.name == val_50_node.name)
        g.node.insert(val_50_pos + 1 + {"q": 0, "k": 1, "v": 2}[suffix], new_concat)

    # --- 3. Reconnect Reshape consumers ---
    # Per layer L (0..27):
    #   linear_{7L+0} (Q weight dims [1024,2048]) -> node_view_{4L+0}
    #   linear_{7L+1} (K weight dims [1024,1024]) -> node_view_{4L+1}
    #   linear_{7L+2} (V weight dims [1024,1024]) -> node_view_{4L+2}
    #
    # So: node_view_{4L+0} uses val_50_q, node_view_{4L+1} uses val_50_k,
    #     node_view_{4L+2} uses val_50_v

    consumers = [n for n in g.node if "val_50" in n.input]
    import re

    q_patches = 0
    k_patches = 0
    v_patches = 0

    for n in consumers:
        name = n.name
        match = re.match(r"node_view_(\d+)$", name)
        if match:
            idx = int(match.group(1))
            # Q: idx % 4 == 0, K: idx % 4 == 1, V: idx % 4 == 2
            if idx % 4 == 0:
                replacement = "val_50_q"
                q_patches += 1
            elif idx % 4 == 1:
                replacement = "val_50_k"
                k_patches += 1
            elif idx % 4 == 2:
                replacement = "val_50_v"
                v_patches += 1
            else:
                print(f"  WARNING: unexpected index {idx} for {name}, skipping")
                continue
        elif name == "node_view":
            # node_view (no suffix) is Q for layer 0
            replacement = "val_50_q"
            q_patches += 1
        else:
            print(f"  WARNING: unexpected consumer {name}, skipping")
            continue

        # Replace val_50 with the specific variant in the Reshape's inputs
        new_inputs = [
            replacement if inp == "val_50" else inp for inp in n.input
        ]
        del n.input[:]
        n.input.extend(new_inputs)

    print(f"  Q reshapes patched: {q_patches}")
    print(f"  K reshapes patched: {k_patches}")
    print(f"  V reshapes patched: {v_patches}")
    print(f"  Total: {q_patches + k_patches + v_patches}")

    # --- 4. Clean up: remove original val_41 and val_50 (optional, but cleaner) ---
    # We keep them to avoid breaking anything else that might reference them.
    # val_50 only has the 84 Reshape consumers; val_41 might be used elsewhere.

    # --- 5. Check op checksums / validate ---
    try:
        onnx.checker.check_model(m)
        print("  ONNX validation: PASSED")
    except Exception as e:
        print(f"  ONNX validation WARNING: {e}")
        print("  (non-critical, DirectML may still accept it)")

    # --- 6. Save ---
    # Handle external data files
    output_dir = os.path.dirname(output_path) or "."
    os.makedirs(output_dir, exist_ok=True)

    # If the model has external data, copy the .data file
    input_dir = os.path.dirname(input_path)
    data_file = input_path + ".data"
    if os.path.exists(data_file):
        output_data = output_path + ".data"
        print(f"  Copying external data: {data_file} -> {output_data}")
        shutil.copy2(data_file, output_data)

    onnx.save(m, output_path)
    size_mb = os.path.getsize(output_path) / (1024 * 1024)
    print(f"  Saved: {output_path} ({size_mb:.1f} MB)")
    print("Done!")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python fix_octen_dml.py <input.onnx> [output.onnx]")
        sys.exit(1)

    input_path = sys.argv[1]
    output_path = sys.argv[2] if len(sys.argv) > 2 else input_path.replace(".onnx", "_dml.onnx")
    patch_octen_for_directml(input_path, output_path)