src/imagenet/matching_utils.py

import os
import jax
import time
import copy
import jax.nn as nn
import jax.lax as lax
import jax.numpy as jnp
from lmc_model import print_model
from collections import defaultdict
from typing import NamedTuple
from flax.core import freeze, unfreeze
from jax import random, tree_util, jit, grad, value_and_grad
from scipy.optimize import linear_sum_assignment, minimize
import numpy as np
import matplotlib.pyplot as plt
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

@jax.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 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):
    key =np.array(attn['attention']['key']['kernel'])
    key_bias = np.array(attn['attention']['key']['bias'])
    query = np.array(attn['attention']['query']['kernel'])
    query_bias = np.array(attn['attention']['query']['bias'])
    value = np.array(attn['attention']['value']['kernel'])
    value_bias = np.array(attn['attention']['value']['bias'])
    out = np.array(attn['output']['dense']['kernel'])
    out_bias = np.array(attn['output']['dense']['bias'])
    return  query, key, value, query_bias, key_bias, value_bias, out, out_bias
def reshape_attention_weights(query, key, value, query_bias, key_bias, value_bias, out_kernel, num_heads):
    D = query.shape[0]
    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] if axis == 0 else tensor[i * D_k:(i + 1) * D_k,:]
            for i in range(num_heads)
        ])
    def stack_bias_per_head(bias):
        return np.stack([bias[i * D_k:(i + 1) * D_k] 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 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, activations, alpha=0.5):
#     # activations: (B, L, D)
#     B, L, D = activations.shape
#     print("Activations shape: ", activations.shape)
#     # Augment activations with a column of ones for bias calculation
#     ones_col = jnp.ones((B, L, 1))
#     X_tilde = jnp.concatenate([activations, ones_col], axis=-1)  # Shape (B, L, D+1)
#     d_head = W_Q_a[0].shape[1]
#     sqrt_d = jnp.sqrt(float(d_head))
#     C = np.zeros((num_heads, num_heads))
#     for i in range(num_heads):
#         # Pre-compute for model A, head i
#         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])
#         Q_a_i = X_tilde @ tilde_W_Q_a_i
#         K_a_i = X_tilde @ tilde_W_K_a_i
#         S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d
#         P_a_i = nn.softmax(S_a_i, axis=-1)
#         V_tilde_a_i = X_tilde @ tilde_W_V_a_i
#         V_a_i = V_tilde_a_i @ W_O_a[i]
#         for j in range(num_heads):
#             # Compute for model B, head j
#             tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j])
#             tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j])
#             tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j])
#             Q_b_j = X_tilde @ tilde_W_Q_b_j
#             K_b_j = X_tilde @ tilde_W_K_b_j
#             S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d
#             P_b_j = nn.softmax(S_b_j, axis=-1)
#             V_tilde_b_j = X_tilde @ tilde_W_V_b_j
#             V_b_j = V_tilde_b_j @ W_O_b[j]
#             cost_P = jnp.sum((P_a_i - P_b_j)**2)
#             cost_V = jnp.sum((V_a_i - V_b_j)**2)
#             total_cost = (alpha * cost_P + (1 - alpha) * cost_V) / B
#             C[i, j] = total_cost
#     return C
# 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,
#                                num_heads, activations, alpha=0.5, epsilon=1e-8):
#     """
#     Computes the cost matrix for attention head permutation using Cosine Similarity.
#     The cost C[i, j] is defined as 1 - cosine_similarity, which is minimized when
#     the output tensors are most similar.
#     """
#     # activations: (B, L, D)
#     B, L, D = activations.shape
#     print("Activations shape: ", activations.shape)
#     # Augment activations with a column of ones for bias calculation
#     ones_col = jnp.ones((B, L, 1))
#     X_tilde = jnp.concatenate([activations, ones_col], axis=-1)  # Shape (B, L, D+1)
#     d_head = W_Q_a[0].shape[1]
#     sqrt_d = jnp.sqrt(float(d_head))
#     C = np.zeros((num_heads, num_heads))
#     for i in range(num_heads):
#         # Pre-compute for model A, head i
#         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])
#         Q_a_i = X_tilde @ tilde_W_Q_a_i
#         K_a_i = X_tilde @ tilde_W_K_a_i
#         S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d
#         P_a_i = nn.softmax(S_a_i, axis=-1)
#         P_a_i_flat = P_a_i.flatten()
#         V_tilde_a_i = X_tilde @ tilde_W_V_a_i
#         V_a_i = V_tilde_a_i @ W_O_a[i]
#         V_a_i_flat = V_a_i.flatten()
#         for j in range(num_heads):
#             # Compute for model B, head j
#             tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j])
#             tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j])
#             tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j])
#             Q_b_j = X_tilde @ tilde_W_Q_b_j
#             K_b_j = X_tilde @ tilde_W_K_b_j
#             S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d
#             P_b_j = nn.softmax(S_b_j, axis=-1)
#             P_b_j_flat = P_b_j.flatten()
#             V_tilde_b_j = X_tilde @ tilde_W_V_b_j
#             V_b_j = V_tilde_b_j @ W_O_b[j]
#             V_b_j_flat = V_b_j.flatten()
#             # Cosine similarity for P
#             dot_P = jnp.dot(P_a_i_flat, P_b_j_flat)
#             norm_P_a = jnp.linalg.norm(P_a_i_flat)
#             norm_P_b = jnp.linalg.norm(P_b_j_flat)
#             cos_sim_P = dot_P / (norm_P_a * norm_P_b + epsilon)
#             cost_P = 1.0 - cos_sim_P
#             # Cosine similarity for V
#             dot_V = jnp.dot(V_a_i_flat, V_b_j_flat)
#             norm_V_a = jnp.linalg.norm(V_a_i_flat)
#             norm_V_b = jnp.linalg.norm(V_b_j_flat)
#             cos_sim_V = dot_V / (norm_V_a * norm_V_b + epsilon)
#             cost_V = 1.0 - cos_sim_V
#             total_cost = (alpha * cost_P + (1 - alpha) * cost_V)
#             C[i, j] = total_cost
#     return C
# Version 3: Normalized Norm
# 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,
#                                    num_heads, activations, alpha=0.5, epsilon=1e-8):
#     """
#     Computes the cost matrix for attention head permutation using Normalized Frobenius Norm.
#     The cost C[i, j] is the relative error: ||A-B||_F / ||A||_F.
#     """
#     # activations: (B, L, D)
#     B, L, D = activations.shape
#     print("Activations shape: ", activations.shape)
#     # Augment activations with a column of ones for bias calculation
#     ones_col = jnp.ones((B, L, 1))
#     X_tilde = jnp.concatenate([activations, ones_col], axis=-1)  # Shape (B, L, D+1)
#     d_head = W_Q_a[0].shape[1]
#     sqrt_d = jnp.sqrt(float(d_head))
#     C = np.zeros((num_heads, num_heads))
#     for i in range(num_heads):
#         # Pre-compute for model A, head i
#         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])
#         Q_a_i = X_tilde @ tilde_W_Q_a_i
#         K_a_i = X_tilde @ tilde_W_K_a_i
#         S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d
#         P_a_i = nn.softmax(S_a_i, axis=-1)
#         V_tilde_a_i = X_tilde @ tilde_W_V_a_i
#         V_a_i = V_tilde_a_i @ W_O_a[i]
#         for j in range(num_heads):
#             # Compute for model B, head j
#             tilde_W_Q_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j])
#             tilde_W_K_b_j = compute_extended_weights(W_K_b[j], b_K_b[j])
#             tilde_W_V_b_j = compute_extended_weights(W_V_b[j], b_V_b[j])
#             Q_b_j = X_tilde @ tilde_W_Q_b_j
#             K_b_j = X_tilde @ tilde_W_K_b_j
#             S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d
#             P_b_j = nn.softmax(S_b_j, axis=-1)
#             V_tilde_b_j = X_tilde @ tilde_W_V_b_j
#             V_b_j = V_tilde_b_j @ W_O_b[j]
#             # Normalized Frobenius norm for P. We flatten the 3D tensor to compute the L2 norm.
#             norm_diff_P = jnp.linalg.norm((P_a_i - P_b_j).flatten())
#             norm_P_a = jnp.linalg.norm(P_a_i.flatten())
#             cost_P = norm_diff_P / (norm_P_a + epsilon)
#             # Normalized Frobenius norm for V. We flatten the 3D tensor to compute the L2 norm.
#             norm_diff_V = jnp.linalg.norm((V_a_i - V_b_j).flatten())
#             norm_V_a = jnp.linalg.norm(V_a_i.flatten())
#             cost_V = norm_diff_V / (norm_V_a + epsilon)
#             total_cost = (alpha * cost_P + (1 - alpha) * cost_V)
#             C[i, j] = total_cost
#     return C
def 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 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 = {}
    for i in range(h):
        print(f"Aligning Heads {i}")
        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_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)])
        A_i_init = solve_orthogonal(tilde_W_Q_a_i, tilde_W_Q_b_i, tilde_W_K_a_i, tilde_W_K_b_i)
        # A_i = optimize(A_i_init, 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
        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)])
        Y_O = W_O_a[i].T
        Y_O_prime = W_O_b[i].T
        B_i_init = solve_orthogonal(tilde_W_V_a_i, tilde_W_V_b_i, Y_O, Y_O_prime)
        # B_i = optimize(B_i_init, tilde_W_V_a_i, tilde_W_V_b_i, Y_O, Y_O_prime)
        B_i = B_i_init
        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]

        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}
        }
    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)
    # print(query_kernel.shape,"|",query_bias.shape)
    # print(key_kernel.shape,"|",key_bias.shape)
    # print(value_kernel.shape,"|",value_bias.shape)
    # print(out_kernel.shape,)
    query_kernel = query_kernel.reshape(-1, D)
    key_kernel = key_kernel.reshape(-1, D)
    value_kernel = value_kernel.reshape(-1, D)
    out_kernel = out_kernel.reshape(D, -1)
    query_bias = query_bias.reshape(-1)
    key_bias = key_bias.reshape(-1)
    value_bias = value_bias.reshape(-1)
    return {
        'attention': {
            'query': {'kernel': jnp.array(query_kernel),'bias': jnp.array(query_bias),},
            'value': {'kernel': jnp.array(value_kernel),'bias': jnp.array(value_bias),},
            'key':  {'kernel': jnp.array(key_kernel),'bias': jnp.array(key_bias),},
        },
        'output': {
            'dense': {'kernel': jnp.array(out_kernel),'bias': jnp.array(out_bias_b),},
        },
    }

def align_attention_params(rng, params_a, params_b, layer_idx, num_heads, activation, permute_heads=True, optimize=False, alpha=0.5):
    attn_a = params_a['vit']['encoder']["layer"][str(layer_idx)]['attention']
    attn_b = params_b['vit']['encoder']["layer"][str(layer_idx)]['attention']
    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 = {}, {}
    for i in range(num_heads):
        result = 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
    return return_dict 

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.lmc_config.num_attention_heads,
                activations_for_layer, permute_heads=permute_heads, optimize=optimize
            )
            aligned_params['vit']['encoder']["layer"][str(layer_idx)]['attention'] = 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.lmc_config.num_attention_heads)
    #     cost_head['vit']['encoder']["layer"][str(layer_idx)]['attention'] = 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.lmc_config.num_attention_heads)
    #     naive_head['vit']['encoder']["layer"][str(layer_idx)]['attention'] = aligned_attention_params
    # return {"cost_head": cost_head, "naive_head": naive_head}