src/lgmodeling/matching_utils.py

import os
import copy
import time
import random
import itertools
import numpy as np
import jax.lax as lax
import jax.numpy as jnp
from utils import rngmix
import matplotlib.pyplot as plt
from typing import NamedTuple
from collections import defaultdict
from flax.core import freeze, unfreeze
from scipy.optimize import linear_sum_assignment, minimize
from jax import random, tree_util, jit, grad, value_and_grad
def compute_objective(A, X, X_prime, Y, Y_prime):
    A_inv = np.linalg.inv(A)
    term1 = X - X_prime @ A.T
    term2 = Y - Y_prime @ A_inv
    return np.sum(term1**2) + np.sum(term2**2)

def compute_gradient(A, X, X_prime, Y, Y_prime):
    A_inv = np.linalg.inv(A)
    term1 = -2 * X.T @ X_prime + 2 * A @ X_prime.T @ X_prime
    term2 = 2 * A_inv.T @ Y_prime.T @ (Y - Y_prime @ A_inv) @ A_inv.T
    return term1 + term2

def line_search(A, grad, X, X_prime, Y, Y_prime, max_step=1, tau=0.5, c1=1e-4):
    eta = max_step
    f_current = compute_objective(A, X, X_prime, Y, Y_prime)
    grad_norm2 = np.sum(grad**2)
    n = A.shape[0]
    while eta > 1e-10:
        A_new = A - eta * grad
        if np.linalg.matrix_rank(A_new) < n:
            eta *= tau
            continue
        f_new = compute_objective(A_new, X, X_prime, Y, Y_prime)
        if f_new <= f_current - c1 * eta * grad_norm2:
            return eta
        eta *= tau
    return 0
@jit
def compute_objective_jax(A, X, X_prime, Y, Y_prime, cond_threshold=1e6):
    cond = jnp.linalg.cond(A)
    def safe_obj():
        A_inv = jnp.linalg.inv(A)
        term1 = X - X_prime @ A.T
        term2 = Y - Y_prime @ A_inv
        return jnp.sum(term1**2) + jnp.sum(term2**2)
    return lax.cond(cond > cond_threshold, lambda: jnp.inf, safe_obj)
compute_value_and_grad_jax = jit(value_and_grad(compute_objective_jax))
def solve_orthogonal(X, X_prime, Y, Y_prime):
    B = X.T @ X_prime + Y.T @ Y_prime
    U, _, Vt = np.linalg.svd(B)
    return U @ Vt
def solve_rope(X, X_prime, Y, Y_prime,max_iters=200, tol=1e-16):
    d = X.shape[1]
    assert d % 2 == 0, "d must be even."
    assert X.shape[1] == X_prime.shape[1] == Y.shape[1] == Y_prime.shape[1]
    def rot(theta):
        c, s = np.cos(theta), np.sin(theta)
        return np.array([[c, -s], [s, c]])
    def block_cols(j): return [2*j, 2*j+1]
    def solve_block(Q1blk, Q2blk, K1blk, K2blk):
        A, B, Ah, Bh = Q1blk, Q2blk, K1blk, K2blk
        a, ah = np.sum(A*A), np.sum(Ah*Ah)
        c_const = np.sum(B*B) + np.sum(Bh*Bh)
        C, Ch = A.T @ B, Ah.T @ Bh
        t_tr, s_sk = np.trace(C), C[0,1] - C[1,0]
        th_tr, sh_sk = np.trace(Ch), Ch[0,1] - Ch[1,0]
        u, v, w = t_tr**2+s_sk**2, th_tr**2+sh_sk**2, t_tr*th_tr+s_sk*sh_sk
        eps = 1e-18
        def phi(t): return max(u*t + v/max(t,eps) + 2*w, 0.0)
        def gprime(t):
            denom = np.sqrt(phi(t))
            if denom < eps: return a - ah/(t*t)
            return (a - ah/(t*t)) - (u - v/(t*t)) / denom
        t0 = np.sqrt((ah+eps)/(a+eps))
        t_lo, gp_lo = t0, gprime(t0)
        if gp_lo < 0.0:
            t_hi = t_lo
            for _ in range(max_iters):
                t_hi *= 2.0
                if gprime(t_hi) >= 0.0: break
        else:
            t_hi = t_lo
            for _ in range(max_iters):
                t_lo *= 0.5
                if gprime(t_lo) <= 0.0: break
        def gval(t): return a*t + ah/max(t,eps) + c_const - 2*np.sqrt(phi(t))
        if not (gprime(t_lo) <= 0.0 <= gprime(t_hi)):
            t_star = min([(t_lo,gval(t_lo)),(t_hi,gval(t_hi))], key=lambda z:z[1])[0]
        else:
            for _ in range(max_iters):
                t_mid = 0.5*(t_lo+t_hi)
                gp_mid = gprime(t_mid)
                if abs(gp_mid) < tol or (t_hi-t_lo) <= tol*(1+t_mid):
                    t_star = t_mid; break
                if gp_mid < 0.0: t_lo = t_mid
                else: t_hi = t_mid
            else:
                t_star = 0.5*(t_lo+t_hi)
        rho = np.sqrt(max(t_star, eps))
        alpha = rho*t_tr + (1/rho)*th_tr
        beta = rho*s_sk + (1/rho)*sh_sk
        theta = np.arctan2(beta, alpha)
        return rho, theta
    P = np.zeros((d, d))
    for j in range(d//2):
        cols = block_cols(j)
        rho, theta = solve_block(X[:,cols], X_prime[:,cols], Y[:,cols], Y_prime[:,cols])
        P[np.ix_(cols, cols)] = rho * rot(theta)
    return P


def optimize_alignment(A_init, X, X_prime, Y, Y_prime, max_iter=5000):
    objective_values = []
    grad_norms = []
    condition_nums = []

    def obj_fn(flat_A):
        A = flat_A.reshape(A_init.shape)
        obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime))
        return float(obj), np.array(grad_val).flatten()

    def callback(flat_A):
        A = flat_A.reshape(A_init.shape)
        obj, grad_val = compute_value_and_grad_jax(jnp.array(A), jnp.array(X), jnp.array(X_prime), jnp.array(Y), jnp.array(Y_prime))
        grad_norm = jnp.linalg.norm(grad_val, 'fro')
        cond = jnp.linalg.cond(jnp.array(A))
        objective_values.append(float(obj))
        grad_norms.append(float(grad_norm))
        condition_nums.append(float(cond))

    res = minimize(obj_fn, A_init.flatten(), jac=True, method='L-BFGS-B', options={'maxiter': max_iter}, callback=callback)
    A_opt = res.x.reshape(A_init.shape)
    return A_opt, objective_values, grad_norms, condition_nums


def extract_attention_params(attn):
    c_attn_kernel = np.array(attn['c_attn']['kernel'])
    c_attn_bias = np.array(attn['c_attn']['bias'])
    c_proj_kernel = np.array(attn['c_proj']['kernel'])
    c_proj_bias = np.array(attn['c_proj']['bias'])
    query, key, value = np.split(c_attn_kernel, 3, axis=0)
    query_bias, key_bias, value_bias = np.split(c_attn_bias, 3, axis=0)
    return query, key, value, query_bias, key_bias, value_bias, c_proj_kernel, c_proj_bias
def reshape_attention_weights(query, key, value, query_bias, key_bias, value_bias, out_kernel, num_heads):
    D = query.shape[1]
    D_k = D_v = D // num_heads
    def stack_per_head(tensor, axis=0):
        return np.stack([
            tensor[i * D_k:(i + 1) * D_k, :].T if axis == 0 else tensor[:, i * D_k:(i + 1) * D_k].T
            for i in range(num_heads)
        ])
    def stack_bias_per_head(bias):
        return np.stack([bias[i * D_k:(i + 1) * D_k].T for i in range(num_heads)])
    W_Q = stack_per_head(query)
    W_K = stack_per_head(key)
    W_V = stack_per_head(value)
    W_O = stack_per_head(out_kernel, axis=1)
    b_Q = stack_bias_per_head(query_bias)
    b_K = stack_bias_per_head(key_bias)
    b_V = stack_bias_per_head(value_bias)
    return W_Q, b_Q, W_K, b_K, W_V, b_V, W_O
def compute_extended_weights(W, b):
    return np.vstack([W, b.reshape(1, -1)])
def compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a,
                        W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b,
                        h, activations, alpha=0.5):
    C = np.zeros((h, h))
    for i in range(h):
        tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)])
        tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)])
        tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)])
        QKT_a_i = tilde_W_Q_a_i @ tilde_W_K_a_i.T
        VO_a_i = tilde_W_V_a_i @ W_O_a[i]
        centered_QKT_a_i = QKT_a_i - np.mean(QKT_a_i, axis=1, keepdims=True)

        for j in range(h):
            tilde_W_Q_b_j = np.vstack([W_Q_b[j], b_Q_b[j].reshape(1, -1)])
            tilde_W_K_b_j = np.vstack([W_K_b[j], b_K_b[j].reshape(1, -1)])
            tilde_W_V_b_j = np.vstack([W_V_b[j], b_V_b[j].reshape(1, -1)])
            QKT_b_j = tilde_W_Q_b_j @ tilde_W_K_b_j.T
            VO_b_j = tilde_W_V_b_j @ W_O_b[j]
            centered_QKT_b_j = QKT_b_j - np.mean(QKT_b_j, axis=1, keepdims=True)

            cost = 0.5 * np.sum((centered_QKT_a_i - centered_QKT_b_j) ** 2)
            cost += 0.5 * np.sum((VO_a_i - VO_b_j) ** 2)
            C[i, j] = cost
    return C
def additive_align_single_head(W_Q_a_i, b_Q_a_i, W_K_a_i, b_K_a_i, W_V_a_i, b_V_a_i, W_O_a_i,
                      W_Q_b_i, b_Q_b_i, W_K_b_i, b_K_b_i, W_V_b_i, b_V_b_i, W_O_b_i, optimize):
    tilde_W_Q_a_i = compute_extended_weights(W_Q_a_i, b_Q_a_i)
    tilde_W_K_a_i = compute_extended_weights(W_K_a_i, b_K_a_i)
    tilde_W_V_a_i = compute_extended_weights(W_V_a_i, b_V_a_i)
    Y_O_a_i = W_O_a_i.T
    tilde_W_Q_b_i = compute_extended_weights(W_Q_b_i, b_Q_b_i)
    tilde_W_K_b_i = compute_extended_weights(W_K_b_i, b_K_b_i)
    tilde_W_V_b_i = compute_extended_weights(W_V_b_i, b_V_b_i)
    Y_O_b_i = W_O_b_i.T
    A_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i)
    B_init = solve_orthogonal(Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i)
    if optimize:
        A, objective_values_A, grad_norms_A, condition_nums_A = optimize_alignment(
            A_init, tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i
        )
        B, objective_values_B, grad_norms_B, condition_nums_B = optimize_alignment(
            B_init, Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i
        )
    else:
        A = A_init
        B = B_init
    A_inv = np.linalg.inv(A)
    B_inv = np.linalg.inv(B)
    W_Q_aligned = W_Q_b_i @ A.T
    b_Q_aligned = b_Q_b_i @ A.T
    W_K_aligned = W_K_b_i @ A_inv
    b_K_aligned = b_K_b_i @ A_inv
    W_V_aligned = W_V_b_i @ B_inv
    b_V_aligned = b_V_b_i @ B_inv
    W_O_aligned = B @ W_O_b_i
    aligned_params = {
        'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned},
        'key': {'kernel': W_K_aligned, 'bias': b_K_aligned},
        'value': {'kernel': W_V_aligned, 'bias': b_V_aligned},
        'out': {'kernel': W_O_aligned}
    }
    if optimize:
        return {
            'aligned_params': aligned_params,
            'metrics_A': {
                'objective_values': objective_values_A,
                'grad_norms': grad_norms_A,
                'condition_nums': condition_nums_A
            },
            'metrics_B': {
                'objective_values': objective_values_B,
                'grad_norms': grad_norms_B,
                'condition_nums': condition_nums_B
            }
        }
    return {'aligned_params': aligned_params}

def _fro2(x):
    if x.ndim == 1:   # vector -> Euclidean norm
        return float(np.linalg.norm(x)**2)
    else:             # matrix -> Frobenius norm
        return float(np.linalg.norm(x, 'fro')**2)

def rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b,
                         W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, h):
    aligned_params = {}

    total_pre = 0.0
    total_post = 0.0
    # --- Pairwise totals ---
    total_qk_pre = 0.0
    total_qk_post = 0.0
    total_vo_pre = 0.0
    total_vo_post = 0.0

    for i in range(h):
        # ===== Build augmented (kernel+bias row) for A (Q,K) =====
        tilde_W_Q_a_i = np.vstack([W_Q_a[i], b_Q_a[i].reshape(1, -1)])
        tilde_W_K_a_i = np.vstack([W_K_a[i], b_K_a[i].reshape(1, -1)])

        # ===== Build augmented (kernel+bias row) for B (Q,K) =====
        tilde_W_Q_b_i = np.vstack([W_Q_b[i], b_Q_b[i].reshape(1, -1)])
        tilde_W_K_b_i = np.vstack([W_K_b[i], b_K_b[i].reshape(1, -1)])

        # ===== Solve A_i for (Q,K) pair =====
        A_i_init = solve_rope(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i)
        A_i = A_i_init  # (optionally run a refinement step)

        # ===== Build augmented (kernel+bias row) for (V,O) pair =====
        tilde_W_V_a_i = np.vstack([W_V_a[i], b_V_a[i].reshape(1, -1)])
        tilde_W_V_b_i = np.vstack([W_V_b[i], b_V_b[i].reshape(1, -1)])

        # For O we transpose and pad with one zero row to match augmented shape
        Y_O       = W_O_a[i].T
        Y_O_prime = W_O_b[i].T
        Y_O_padded        = np.vstack([Y_O,        np.zeros((1, Y_O.shape[1]))])
        Y_O_prime_padded  = np.vstack([Y_O_prime,  np.zeros((1, Y_O_prime.shape[1]))])

        # ===== Solve B_i for (V,O) pair =====
        B_i_init = solve_rope(tilde_W_V_a_i, tilde_W_V_b_i, Y_O_padded, Y_O_prime_padded)
        B_i = B_i_init  # (optionally run a refinement step)

        # ===== Apply transforms =====
        A_i_inv = np.linalg.inv(A_i)
        B_i_inv = np.linalg.inv(B_i)

        # W_Q_aligned = W_Q_b[i] @ A_i.T
        # b_Q_aligned = b_Q_b[i] @ A_i.T

        # W_K_aligned = W_K_b[i] @ A_i_inv
        # b_K_aligned = b_K_b[i] @ A_i_inv

        # W_V_aligned = W_V_b[i] @ B_i_inv
        # b_V_aligned = b_V_b[i] @ B_i_inv

        # W_O_aligned = B_i @ W_O_b[i]
        W_Q_aligned = W_Q_b[i] @ A_i_inv.T     # was A_i.T   -> FIX: A_i^{-T}
        b_Q_aligned = b_Q_b[i] @ A_i_inv.T

        W_K_aligned = W_K_b[i] @ A_i           # was A_i_inv -> FIX: A_i
        b_K_aligned = b_K_b[i] @ A_i

        # --- V,O pair: use P for V (right-multiply), and P^{-1} for O (left-multiply) ---
        W_V_aligned = W_V_b[i] @ B_i           # was B_i_inv -> FIX: B_i
        b_V_aligned = b_V_b[i] @ B_i

        W_O_aligned = B_i_inv @ W_O_b[i]       # was B_i @ W_O_b[i] -> FIX: B_i^{-1} on the left
        aligned_params[f'head_{i}'] = {
            'query': {'kernel': W_Q_aligned, 'bias': b_Q_aligned},
            'key':   {'kernel': W_K_aligned, 'bias': b_K_aligned},
            'value': {'kernel': W_V_aligned, 'bias': b_V_aligned},
            'out':   {'kernel': W_O_aligned}
        }

        # ===== Frobenius^2 BEFORE (a vs raw b) =====
        pre_q = _fro2(W_Q_a[i] - W_Q_b[i]) + _fro2(b_Q_a[i] - b_Q_b[i])
        pre_k = _fro2(W_K_a[i] - W_K_b[i]) + _fro2(b_K_a[i] - b_K_b[i])
        pre_v = _fro2(W_V_a[i] - W_V_b[i]) + _fro2(b_V_a[i] - b_V_b[i])
        pre_o = _fro2(W_O_a[i] - W_O_b[i])  # O has no bias in your structure

        # ===== Frobenius^2 AFTER (a vs aligned b) =====
        post_q = _fro2(W_Q_a[i] - W_Q_aligned) + _fro2(b_Q_a[i] - b_Q_aligned)
        post_k = _fro2(W_K_a[i] - W_K_aligned) + _fro2(b_K_a[i] - b_K_aligned)
        post_v = _fro2(W_V_a[i] - W_V_aligned) + _fro2(b_V_a[i] - b_V_aligned)
        post_o = _fro2(W_O_a[i] - W_O_aligned)

        pre_sum  = pre_q + pre_k + pre_v + pre_o
        post_sum = post_q + post_k + post_v + post_o

        total_pre  += pre_sum
        total_post += post_sum

        # ===== Pairwise sums =====
        pre_qk  = pre_q + pre_k
        post_qk = post_q + post_k
        pre_vo  = pre_v + pre_o
        post_vo = post_v + post_o

        total_qk_pre  += pre_qk
        total_qk_post += post_qk
        total_vo_pre  += pre_vo
        total_vo_post += post_vo

        # ===== Per-head print =====
        print(f"[Head {i}] Fro^2 pre={pre_sum:.6f}  post={post_sum:.6f} improve={pre_sum - post_sum:.6f}")
        print(f"  Q:  pre={pre_q:.6f} post={post_q:.6f}")
        print(f"  K:  pre={pre_k:.6f} post={post_k:.6f}")
        print(f"  V:  pre={pre_v:.6f} post={post_v:.6f}")
        print(f"  O:  pre={pre_o:.6f} post={post_o:.6f}")
        # --- New: pairwise breakdowns ---
        print(f"  [Q,K] pair: pre={pre_qk:.6f} post={post_qk:.6f} improve={pre_qk - post_qk:.6f}")
        print(f"  [V,O] pair: pre={pre_vo:.6f} post={post_vo:.6f} improve={pre_vo - post_vo:.6f}")
    # ===== Totals =====
    print("=== Frobenius^2 (including biases) ===")
    print(f"Total pre : {total_pre:.6f}")
    print(f"Total post: {total_post:.6f}")
    print(f"Total improvement: {total_pre - total_post:.6f} ({0.0 if total_pre==0 else 100.0*(total_pre-total_post)/total_pre:.2f}%)")

    # --- New: Pairwise totals ---
    print("=== Pairwise Frobenius^2 (including biases) ===")
    print(f"[Q,K] total pre : {total_qk_pre:.6f}")
    print(f"[Q,K] total post: {total_qk_post:.6f}")
    print(f"[Q,K] improvement: {total_qk_pre - total_qk_post:.6f} ({0.0 if total_qk_pre==0 else 100.0*(total_qk_pre-total_qk_post)/total_qk_pre:.2f}%)")
    print(f"[V, O] total pre : {total_vo_pre:.6f}")
    print(f"[V, O] total post: {total_vo_post:.6f}")
    print(f"[V, O] improvement: {total_vo_pre - total_vo_post:.6f} ({0.0 if total_vo_pre==0 else 100.0*(total_vo_pre-total_vo_post)/total_vo_pre:.2f}%)")

    return aligned_params
def merge_aligned_params(aligned_params, h, D, out_bias_b):
    query_kernel = np.stack([aligned_params[f'head_{i}']['query']['kernel'] for i in range(h)], axis=1)
    query_bias = np.stack([aligned_params[f'head_{i}']['query']['bias'] for i in range(h)], axis=0)
    key_kernel = np.stack([aligned_params[f'head_{i}']['key']['kernel'] for i in range(h)], axis=1)
    key_bias = np.stack([aligned_params[f'head_{i}']['key']['bias'] for i in range(h)], axis=0)
    value_kernel = np.stack([aligned_params[f'head_{i}']['value']['kernel'] for i in range(h)], axis=1)
    value_bias = np.stack([aligned_params[f'head_{i}']['value']['bias'] for i in range(h)], axis=0)
    out_kernel = np.stack([aligned_params[f'head_{i}']['out']['kernel'] for i in range(h)], axis=0)

    query_kernel = query_kernel.transpose(1, 2, 0).reshape(-1, D)
    key_kernel = key_kernel.transpose(1, 2, 0).reshape(-1, D)
    value_kernel = value_kernel.transpose(1, 2, 0).reshape(-1, D)
    out_kernel = out_kernel.transpose(2, 0, 1).reshape(D, -1)

    query_bias = query_bias.reshape(-1)
    key_bias = key_bias.reshape(-1)
    value_bias = value_bias.reshape(-1)

    return {
        'c_attn': {
            'kernel': jnp.array(np.concatenate([query_kernel, key_kernel, value_kernel], axis=0)),
            'bias': jnp.array(np.concatenate([query_bias, key_bias, value_bias], axis=0)),
        },
        'c_proj': {'kernel': jnp.array(out_kernel),'bias': jnp.array(out_bias_b),}
    }

def align_attention_params(rng, params_a, params_b, layer_idx, config, activation, permute_heads=True, optimize=False, alpha=0.5):
    num_heads = config.lmc_config.n_head
    attn_a = params_a['transformer']['h'][str(layer_idx)]['attn']
    attn_b = params_b['transformer']['h'][str(layer_idx)]['attn']
    query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a)
    query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b)
    W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads)
    W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads)
    if permute_heads:
        C = compute_cost_matrix(W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a,
                                W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b, num_heads, activation, alpha)
        row_ind, col_ind = linear_sum_assignment(C)
        print("Best Permutation Heads:", col_ind)
        W_Q_b = [W_Q_b[j] for j in col_ind]
        b_Q_b = [b_Q_b[j] for j in col_ind]
        W_K_b = [W_K_b[j] for j in col_ind]
        b_K_b = [b_K_b[j] for j in col_ind]
        W_V_b = [W_V_b[j] for j in col_ind]
        b_V_b = [b_V_b[j] for j in col_ind]
        W_O_b = [W_O_b[j] for j in col_ind]
    if optimize:
        metrics_A_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']}
        metrics_B_all = {key: [] for key in ['objective_values', 'grad_norms', 'condition_nums']}
    aligned_params, return_dict = {}, {}
    if(config.position_embeddings in ["learnable","sinusoidal"]):
        for i in range(num_heads):
            result = additive_align_single_head(
                W_Q_a[i], b_Q_a[i], W_K_a[i], b_K_a[i], W_V_a[i], b_V_a[i], W_O_a[i],
                W_Q_b[i], b_Q_b[i], W_K_b[i], b_K_b[i], W_V_b[i], b_V_b[i], W_O_b[i], optimize
            )
            aligned_params[f'head_{i}'] = result['aligned_params']
            if optimize:
                for key in metrics_A_all:
                    metrics_A_all[key].append(result['metrics_A'][key])
                    metrics_B_all[key].append(result['metrics_B'][key])
        return_dict['aligned_params'] = merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b)
        if optimize:
            return_dict['metrics_A_all'] = metrics_A_all
            return_dict['metrics_B_all'] = metrics_B_all
    elif(config.position_embeddings in ["rope"]): 
        aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b,
                                W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads)
    return return_dict
def permute_align_attention_params(rng, params_a, params_b, layer_idx, config,col_ind):
    num_heads = config.lmc_config.n_head
    attn_a = params_a['transformer']['h'][layer_idx]['attn']
    attn_b = params_b['transformer']['h'][layer_idx]['attn']
    query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a)
    query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b)
    W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads)
    W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads)
    W_Q_b = [W_Q_b[j] for j in col_ind]
    b_Q_b = [b_Q_b[j] for j in col_ind]
    W_K_b = [W_K_b[j] for j in col_ind]
    b_K_b = [b_K_b[j] for j in col_ind]
    W_V_b = [W_V_b[j] for j in col_ind]
    b_V_b = [b_V_b[j] for j in col_ind]
    W_O_b = [W_O_b[j] for j in col_ind]
    if(config.position_embeddings in ["learnable","sinusoidal"]):
        aligned_params = additive_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b,
                                        W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads)
    elif(config.position_embeddings in ["rope"]): 
        aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b,
                                W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads)
    return merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b)
def naive_align_attention_params(rng, params_a, params_b, layer_idx, config):
    num_heads = config.lmc_config.n_head
    attn_a = params_a['transformer']['h'][layer_idx]['attn']
    attn_b = params_b['transformer']['h'][layer_idx]['attn']
    query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, out_bias_a = extract_attention_params(attn_a)
    query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, out_bias_b = extract_attention_params(attn_b)
    W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_attention_weights(query_a, key_a, value_a, query_bias_a, key_bias_a, value_bias_a, out_a, num_heads)
    W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_attention_weights(query_b, key_b, value_b, query_bias_b, key_bias_b, value_bias_b, out_b, num_heads)
    if(config.position_embeddings in ["learnable","sinusoidal"]):
        aligned_params = additive_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b,
                                        W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads)
    elif(config.position_embeddings in ["rope"]): 
        aligned_params = rope_apply_alignment(W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b,
                                W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a, num_heads)
    return merge_aligned_params(aligned_params, num_heads, query_a.shape[1], out_bias_b)
def all_matching_attn(rng, params_a, params_b, config):
    results = {}
    permutations = list(itertools.permutations(range(config.lmc_config.n_head)))
    if config.lmc_config.n_head > 4:
        permutations = random.sample(permutations, 24)
    for perm in permutations:
        print("Permutation",perm)
        temp_params = copy.deepcopy(params_b)
        for layer_idx in config.lmc_layer_indices:
            aligned_attention_params = permute_align_attention_params(rng, params_a, params_b, str(layer_idx), config, perm)
            temp_params['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params
        results[str(perm)] = temp_params
    return results

def weight_matching_attn(rng, params_a, params_b, activation, config):
    params_dict = {}
    configurations = [
        ("permu_head_init_ortho_no_opt", 'ortho', True, False),
        ("permu_head_init_ortho_opt", 'ortho', True, True),
        # ("naive_head_init_ortho_no_opt", 'ortho', False, False),
        # ("naive_head_init_ortho_opt", 'ortho', False, True),
    ]
    for name, init_method, permute_heads, optimize in configurations:
        aligned_params = copy.deepcopy(params_b)
        if optimize:
            layer_to_metrics_A = {}
            layer_to_metrics_B = {}
        for layer_idx in config.lmc_layer_indices:
            if activation is not None: activations_for_layer = activation[layer_idx]
            else: activations_for_layer = None
            result = align_attention_params(
                rng, params_a, aligned_params, layer_idx, config,
                activations_for_layer, permute_heads=permute_heads, optimize=optimize
            )
            aligned_params['transformer']['h'][str(layer_idx)]['attn'] = result['aligned_params']
            if optimize:
                layer_to_metrics_A[layer_idx] = result['metrics_A_all']
                layer_to_metrics_B[layer_idx] = result['metrics_B_all']
        
        total_sum = tree_util.tree_reduce(lambda acc, x: acc + jnp.sum(x), aligned_params, initializer=0)
        print(f"{name}: {total_sum}, sanity check")
        params_dict[name] = aligned_params
    return params_dict

    # cost_head = copy.deepcopy(params_b)
    # naive_head = copy.deepcopy(params_b)
    # for layer_idx in config.lmc_layer_indices:
    #     aligned_attention_params = cost_align_attention_params(rng, params_a, params_b, str(layer_idx), config)
    #     cost_head['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params
    # for layer_idx in config.lmc_layer_indices:
    #     aligned_attention_params = naive_align_attention_params(rng, params_a, params_b, str(layer_idx), config)
    #     naive_head['transformer']['h'][str(layer_idx)]['attn'] = aligned_attention_params
    # return {"cost_head": cost_head, "naive_head": naive_head}