#!/usr/bin/env python3
"""Operator-fusion at the weight level: merge FFN w1 (gate) and w3 (up) into one w13 tensor.

For a Linear weight stored as [in, out] (ggml ne = [in, out]), out-rows are contiguous blocks,
so concatenating raw quantized bytes along the output dimension yields a valid quantized
tensor of [in, out1 + out2]. Output columns [0, out1) equal the original w1 output and
[out1, out1 + out2) equal w3 output, bit-for-bit, because each output element is the same dot product.
"""
import sys
import gguf
import numpy as np

src, dst = sys.argv[1], sys.argv[2]
reader = gguf.GGUFReader(src)
tensors = {t.name: t for t in reader.tensors}

prefixes = sorted({n[: -len("feed_forward.w1.weight")] for n in tensors if n.endswith("feed_forward.w1.weight")})
fused_names = {}
for p in prefixes:
    w1, w3 = tensors[p + "feed_forward.w1.weight"], tensors[p + "feed_forward.w3.weight"]
    assert w1.tensor_type == w3.tensor_type, p
    assert list(w1.shape) == list(w3.shape), p
    fused_names[p + "feed_forward.w1.weight"] = p + "feed_forward.w13.weight"
    fused_names[p + "feed_forward.w3.weight"] = None  # consumed

writer = gguf.GGUFWriter(dst, arch="")
n_fused = 0
for t in reader.tensors:
    if t.name in fused_names and fused_names[t.name] is None:
        continue
    if t.name in fused_names:
        p = t.name[: -len("feed_forward.w1.weight")]
        w1 = tensors[p + "feed_forward.w1.weight"]
        w3 = tensors[p + "feed_forward.w3.weight"]
        data = np.concatenate([np.asarray(w1.data), np.asarray(w3.data)], axis=0)  # [out1+out2, row_bytes]
        writer.add_tensor(fused_names[t.name], data, raw_dtype=t.tensor_type)
        n_fused += 1
        continue
    writer.add_tensor(t.name, np.asarray(t.data), raw_dtype=t.tensor_type)

writer.write_header_to_file()
writer.write_kv_data_to_file()
writer.write_tensors_to_file()
writer.close()
print(f"fused {n_fused} FFN pairs -> {dst}")