src/matching_utils.py

from collections import defaultdict
from typing import NamedTuple
from flax.core import freeze, unfreeze
import jax.numpy as jnp
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
import time
import os
import copy
import jax
import jax.lax as lax
import jax.nn as nn
import jax

'''num_heads = 4
    #print all layers
    layer_paths = []
    def collect_layer_paths(path, value):
        # Convert path to a readable string by joining path keys
        path_str = '/'.join([str(p.key) for p in path])
        shape = value.shape
        layer_paths.append((path_str, shape))

    jax.tree_util.tree_map_with_path(collect_layer_paths, pretrained_params)
    print("Layers of the model:")
    for path, shape in layer_paths:  
        print(f"  {path} {shape}")'''
'''example output
    Layers of the model:
    Conv_0/bias (32,)
    Conv_0/kernel (4, 4, 3, 32)
    Dense_0/bias (10,)
    Dense_0/kernel (32, 10)
    TransformerEncoderLayer_0/Dense_0/bias (128,)
    TransformerEncoderLayer_0/Dense_0/kernel (32, 128)
    TransformerEncoderLayer_0/Dense_1/bias (32,)
    TransformerEncoderLayer_0/Dense_1/kernel (128, 32)
    TransformerEncoderLayer_0/LayerNorm_0/bias (32,)
    TransformerEncoderLayer_0/LayerNorm_0/scale (32,)
    TransformerEncoderLayer_0/LayerNorm_1/bias (32,)
    TransformerEncoderLayer_0/LayerNorm_1/scale (32,)
    TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/key/bias (4, 8)
    TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/key/kernel (32, 4, 8) #Key projections for 4 attention heads, each with a dimension of 8 (4 heads x 8 = 32, matching the model's hidden size).
    TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/out/bias (32,)
    TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/out/kernel (4, 8, 32)
    TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/query/bias (4, 8)
    TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/query/kernel (32, 4, 8)
    TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/value/bias (4, 8)
    TransformerEncoderLayer_0/MultiHeadDotProductAttention_0/value/kernel (32, 4, 8)
    TransformerEncoderLayer_1/Dense_0/bias (128,)
    TransformerEncoderLayer_1/Dense_0/kernel (32, 128)
    TransformerEncoderLayer_1/Dense_1/bias (32,)
    TransformerEncoderLayer_1/Dense_1/kernel (128, 32)
    TransformerEncoderLayer_1/LayerNorm_0/bias (32,)
    TransformerEncoderLayer_1/LayerNorm_0/scale (32,)
    TransformerEncoderLayer_1/LayerNorm_1/bias (32,)
    TransformerEncoderLayer_1/LayerNorm_1/scale (32,)
    TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/key/bias (4, 8)
    TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/key/kernel (32, 4, 8)
    TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/out/bias (32,)
    TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/out/kernel (4, 8, 32)
    TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/query/bias (4, 8)
    TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/query/kernel (32, 4, 8)
    TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/value/bias (4, 8)
    TransformerEncoderLayer_1/MultiHeadDotProductAttention_0/value/kernel (32, 4, 8)
    TransformerEncoderLayer_2/Dense_0/bias (128,)
    TransformerEncoderLayer_2/Dense_0/kernel (32, 128)
    TransformerEncoderLayer_2/Dense_1/bias (32,)
    TransformerEncoderLayer_2/Dense_1/kernel (128, 32)
    TransformerEncoderLayer_2/LayerNorm_0/bias (32,)
    TransformerEncoderLayer_2/LayerNorm_0/scale (32,)
    TransformerEncoderLayer_2/LayerNorm_1/bias (32,)
    TransformerEncoderLayer_2/LayerNorm_1/scale (32,)
    TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/key/bias (4, 8)
    TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/key/kernel (32, 4, 8)
    TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/out/bias (32,)
    TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/out/kernel (4, 8, 32)
    TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/query/bias (4, 8)
    TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/query/kernel (32, 4, 8)
    TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/value/bias (4, 8)
    TransformerEncoderLayer_2/MultiHeadDotProductAttention_0/value/kernel (32, 4, 8)
    cls_token (1, 1, 32)
    pos_embedding (1, 65, 32)
'''

def to_numpy(x):
    if isinstance(x, jnp.ndarray):
        return np.array(x)
    return np.array(x)

@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 get_nested_item(d, keys):
    """Accesses a nested dictionary item using a tuple of keys."""
    for key in keys:
        d = d[key]
    return d

# Helper functions for parameter extraction and reshaping
def extract_attention_params(params, layer_idx):
    """
    Extracts MHA parameters from different, known model structures in a compatible way.

    This function detects the model type and constructs the correct path to the
    attention parameters for a given layer index.

    Args:
        params: The Flax parameter tree.
        layer_idx: The integer index of the transformer layer.

    Returns:
        A tuple containing:
        - A flat tuple of the MHA tensors: (key, key_bias, query, query_bias, value, value_bias, out, out_bias).
        - A tuple representing the nested path to the MHA block, for use in updates.
    """
    # Detect vit-jax style: params['Transformer']['encoderblock_...']
    if 'Transformer' in params and f'encoderblock_{layer_idx}' in params['Transformer']:
        mha_path = ('Transformer', f'encoderblock_{layer_idx}', 'MultiHeadDotProductAttention_0')
    # Detect cifar_vit style: params['TransformerEncoderLayer_...']
    elif f'TransformerEncoderLayer_{layer_idx}' in params:
        mha_path = (f'TransformerEncoderLayer_{layer_idx}', 'MultiHeadDotProductAttention_0')
    else:
        raise KeyError(f"Could not find a known path for attention layer {layer_idx} in the provided params.")

    attention_block = get_nested_item(params, mha_path)

    key_k, key_b = attention_block['key']['kernel'], attention_block['key']['bias']
    query_k, query_b = attention_block['query']['kernel'], attention_block['query']['bias']
    value_k, value_b = attention_block['value']['kernel'], attention_block['value']['bias']
    out_k, out_b = attention_block['out']['kernel'], attention_block['out']['bias']
    
    return (key_k, key_b, query_k, query_b, value_k, value_b, out_k, out_b), mha_path

def reshape_to_per_head(params, num_heads):
    """
    Reshapes batched attention parameters into a list of per-head parameters.

    This function assumes a specific shape convention for the input weight and
    bias tensors, which is common in Flax/Linen implementations.

    Args:
        params (dict): A dictionary containing the attention parameters.
            Expected keys and tensor shapes are:
            - 'query': Weight tensor of shape (D, num_heads, d_k)
            - 'query_bias': Bias tensor of shape (num_heads, d_k)
            - 'key': Weight tensor of shape (D, num_heads, d_k)
            - 'key_bias': Bias tensor of shape (num_heads, d_k)
            - 'value': Weight tensor of shape (D, num_heads, d_v)
            - 'value_bias': Bias tensor of shape (num_heads, d_v)
            - 'out': Weight tensor of shape (num_heads, d_v, D)
        num_heads (int): The number of attention heads.

    Returns:
        A tuple containing lists of per-head parameters:
        (W_Q, b_Q, W_K, b_K, W_V, b_V, W_O)
    """

    query_kernel = params['query']
    assert query_kernel.ndim == 3, f"Expected query weights to be 3D, but got shape {query_kernel.shape}"
    assert query_kernel.shape[1] == num_heads, (
        f"The second dimension of the query weight tensor should be num_heads ({num_heads}), "
        f"but got shape {query_kernel.shape}. Please verify your model's parameter shape convention."
    )

    W_Q = [params['query'][:, i, :] for i in range(num_heads)]
    b_Q = [params['query_bias'][i, :] for i in range(num_heads)]
    W_K = [params['key'][:, i, :] for i in range(num_heads)]
    b_K = [params['key_bias'][i, :] for i in range(num_heads)]
    W_V = [params['value'][:, i, :] for i in range(num_heads)]
    b_V = [params['value_bias'][i, :] for i in range(num_heads)]
    W_O = [params['out'][i, :, :] for i in range(num_heads)]
    
    return W_Q, b_Q, W_K, b_K, W_V, b_V, W_O

def compute_extended_weights(W, b):
    return jnp.vstack([jnp.array(W), jnp.array(b).reshape(1, -1)])

# Helper function to plot multiple curves
def plot_multiple_curves(data_list, title, xlabel, ylabel, labels, save_path):
    plt.figure()
    for data, label in zip(data_list, labels):
        label = f"{label} ({data[-1]:.4f})"
        plt.plot(data, label=label)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
    plt.savefig(save_path, bbox_inches='tight')
    plt.close()

# Stage 1 Function: Find Heads Permutation (Data-Dependent)
# Version 1: Using post-softmax probabilities
def compute_cost_matrix_postsoftmax(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_a, activations_b, alpha=0.5, epsilon=1e-8):
    """
    Computes the cost matrix using post-softmax probabilities and model-specific activations.
    """
    B_a, L_a, D_a = activations_a.shape
    B_b, L_b, D_b = activations_b.shape
    assert B_a == B_b and L_a == L_b and D_a == D_b, "Activations for both models must have the same shape."

    # Augment activations for each model separately
    X_tilde_a = jnp.concatenate([activations_a, jnp.ones((B_a, L_a, 1))], axis=-1)
    X_tilde_b = jnp.concatenate([activations_b, jnp.ones((B_b, L_b, 1))], axis=-1)

    d_head = W_Q_a[0].shape[1]
    sqrt_d = jnp.sqrt(float(d_head))
    C = np.zeros((num_heads, num_heads))

    # Pre-compute flattened outputs for model A
    P_flat_a, V_flat_a = [], []
    for i in range(num_heads):
        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_a @ tilde_W_Q_a_i
        K_a_i = X_tilde_a @ 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_flat_a.append(P_a_i.flatten())
        V_a_i = (X_tilde_a @ tilde_W_V_a_i) @ W_O_a[i]
        V_flat_a.append(V_a_i.flatten())

    # Pre-compute flattened outputs for model B
    P_flat_b, V_flat_b = [], []
    for j in range(num_heads):
        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_b @ tilde_W_Q_b_j
        K_b_j = X_tilde_b @ 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_flat_b.append(P_b_j.flatten())
        V_b_j = (X_tilde_b @ tilde_W_V_b_j) @ W_O_b[j]
        V_flat_b.append(V_b_j.flatten())

    # Compute cost matrix from pre-computed values
    for i in range(num_heads):
        for j in range(num_heads):
            # Cosine similarity for P (post-softmax probabilities)
            dot_P = jnp.dot(P_flat_a[i], P_flat_b[j])
            norm_P_a = jnp.linalg.norm(P_flat_a[i])
            norm_P_b = jnp.linalg.norm(P_flat_b[j])
            cost_P = 1.0 - (dot_P / (norm_P_a * norm_P_b + epsilon))

            # Cosine similarity for V (value-projections)
            dot_V = jnp.dot(V_flat_a[i], V_flat_b[j])
            norm_V_a = jnp.linalg.norm(V_flat_a[i])
            norm_V_b = jnp.linalg.norm(V_flat_b[j])
            cost_V = 1.0 - (dot_V / (norm_V_a * norm_V_b + epsilon))
            
            C[i, j] = alpha * cost_P + (1 - alpha) * cost_V

    return C

# Version 2: Using pre-softmax scores
def compute_cost_matrix_presoftmax(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_a, activations_b, alpha=0.5, epsilon=1e-8):
    """
    Computes the cost matrix for attention head permutation using model-specific activations.
    """
    B_a, L_a, D_a = activations_a.shape
    B_b, L_b, D_b = activations_b.shape
    assert B_a == B_b and L_a == L_b and D_a == D_b, "Activations for both models must have the same shape."
    
    # Augment activations for model A
    ones_col_a = jnp.ones((B_a, L_a, 1))
    X_tilde_a = jnp.concatenate([activations_a, ones_col_a], axis=-1)

    # Augment activations for model B
    ones_col_b = jnp.ones((B_b, L_b, 1))
    X_tilde_b = jnp.concatenate([activations_b, ones_col_b], axis=-1)

    d_head = W_Q_a[0].shape[1]
    sqrt_d = jnp.sqrt(float(d_head))
    C = np.zeros((num_heads, num_heads))

    # Pre-compute all head outputs for model A
    S_bar_flat_a, V_flat_a = [], []
    for i in range(num_heads):
        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_a @ tilde_W_Q_a_i
        K_a_i = X_tilde_a @ tilde_W_K_a_i
        S_a_i = jnp.einsum('bld,bmd->blm', Q_a_i, K_a_i) / sqrt_d
        S_bar_a_i = S_a_i - jnp.mean(S_a_i, axis=2, keepdims=True)
        S_bar_flat_a.append(S_bar_a_i.flatten())
        V_tilde_a_i = X_tilde_a @ tilde_W_V_a_i
        V_a_i = V_tilde_a_i @ W_O_a[i]
        V_flat_a.append(V_a_i.flatten())

    # Pre-compute all head outputs for model B
    S_bar_flat_b, V_flat_b = [], []
    for j in range(num_heads):
        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_b @ tilde_W_Q_b_j
        K_b_j = X_tilde_b @ tilde_W_K_b_j
        S_b_j = jnp.einsum('bld,bmd->blm', Q_b_j, K_b_j) / sqrt_d
        S_bar_b_j = S_b_j - jnp.mean(S_b_j, axis=2, keepdims=True)
        S_bar_flat_b.append(S_bar_b_j.flatten())
        V_tilde_b_j = X_tilde_b @ tilde_W_V_b_j
        V_b_j = V_tilde_b_j @ W_O_b[j]
        V_flat_b.append(V_b_j.flatten())

    # Compute cost matrix from pre-computed values
    for i in range(num_heads):
        for j in range(num_heads):
            # Cosine similarity for S
            dot_S = jnp.dot(S_bar_flat_a[i], S_bar_flat_b[j])
            norm_S_a = jnp.linalg.norm(S_bar_flat_a[i])
            norm_S_b = jnp.linalg.norm(S_bar_flat_b[j])
            cos_sim_S = dot_S / (norm_S_a * norm_S_b + epsilon)
            cost_S = 1.0 - cos_sim_S
            # Cosine similarity for V
            dot_V = jnp.dot(V_flat_a[i], V_flat_b[j])
            norm_V_a = jnp.linalg.norm(V_flat_a[i])
            norm_V_b = jnp.linalg.norm(V_flat_b[j])
            cos_sim_V = dot_V / (norm_V_a * norm_V_b + epsilon)
            cost_V = 1.0 - cos_sim_V
            C[i, j] = (alpha * cost_S + (1 - alpha) * cost_V)
    return C

def compute_cost_matrix_data_independent(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, alpha=0.5):
    C = np.zeros((num_heads, num_heads))
    for i in range(num_heads):
        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])
        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(num_heads):
            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])
            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 = alpha * np.sum((centered_QKT_a_i - centered_QKT_b_j)**2) + (1 - alpha) * np.sum((VO_a_i - VO_b_j)**2)
            C[i, j] = cost
    return C

def find_heads_permutation(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_a, activations_b, alpha, data_independent):
    if data_independent:
        C = compute_cost_matrix_data_independent(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, alpha=alpha)
    else:
        C = compute_cost_matrix_presoftmax(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_a, activations_b, alpha)    
    row_ind, col_ind = linear_sum_assignment(to_numpy(C))
    print({int(i): int(j) for i, j in zip(row_ind, col_ind)})
    return row_ind, col_ind

# Stage 2 Function: Align Single Head (with optional optimization)
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,
                      init_method, 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
    
    if init_method == 'ortho':
        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)
    elif init_method == 'random':
        while True:
            A_init = np.random.normal(loc=1, scale=1, size=(tilde_W_Q_a_i.shape[1], tilde_W_Q_a_i.shape[1]))
            if np.linalg.det(A_init) != 0:
                break
        while True:
            B_init = np.random.normal(loc=1, scale=1, size=(tilde_W_V_a_i.shape[1], tilde_W_V_a_i.shape[1]))
            if np.linalg.det(B_init) != 0:
                break
    elif init_method == 'identity':
        A_init = np.eye(tilde_W_Q_a_i.shape[1])
        B_init = np.eye(tilde_W_V_a_i.shape[1])
    else:
        raise ValueError("Invalid initialization method")
    
    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 align_attention_params_main(rng, params_a, params_b, layer_idx, num_heads, 
                                activations_for_layer_a, activations_for_layer_b, plot_path=None, init_method='ortho', permute_heads=True, optimize=True, method_name="", alpha=0.5, data_independent=False):

    params_a_extracted, _ = extract_attention_params(params_a, layer_idx)
    params_b_extracted, mha_path_b = extract_attention_params(params_b, layer_idx)
    
    params_a_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 
                                                   'value', 'value_bias', 'out', 'out_bias'], params_a_extracted)}
    params_b_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 
                                                   'value', 'value_bias', 'out', 'out_bias'], params_b_extracted)}
    
    W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_to_per_head(params_a_np, num_heads)
    W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_to_per_head(params_b_np, num_heads)
    
    if permute_heads:
        row_ind, col_ind = find_heads_permutation(
            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_for_layer_a, activations_for_layer_b, alpha, data_independent
        )
    
        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 = {}
    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],
            init_method, 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])
    
    query_kernel = np.stack([aligned_params[f'head_{i}']['query']['kernel'] for i in range(num_heads)], axis=1)
    query_bias = np.stack([aligned_params[f'head_{i}']['query']['bias'] for i in range(num_heads)], axis=0)
    key_kernel = np.stack([aligned_params[f'head_{i}']['key']['kernel'] for i in range(num_heads)], axis=1)
    key_bias = np.stack([aligned_params[f'head_{i}']['key']['bias'] for i in range(num_heads)], axis=0)
    value_kernel = np.stack([aligned_params[f'head_{i}']['value']['kernel'] for i in range(num_heads)], axis=1)
    value_bias = np.stack([aligned_params[f'head_{i}']['value']['bias'] for i in range(num_heads)], axis=0)
    out_kernel = np.stack([aligned_params[f'head_{i}']['out']['kernel'] for i in range(num_heads)], axis=0)

    return_dict = {
        'aligned_params': {
            'query': {'kernel': jnp.array(query_kernel), 'bias': jnp.array(query_bias)},
            'key': {'kernel': jnp.array(key_kernel), 'bias': jnp.array(key_bias)},
            'value': {'kernel': jnp.array(value_kernel), 'bias': jnp.array(value_bias)},
            'out': {'kernel': jnp.array(out_kernel), 'bias': params_b_np['out_bias']}
        },
        'mha_path': mha_path_b # Return the path for updating model_b
    }
    if optimize:
        return_dict['metrics_A_all'] = metrics_A_all
        return_dict['metrics_B_all'] = metrics_B_all
    return return_dict

def matching_attn(rng, params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads, plot_path):
    params_dict = {}
    configurations = [
        ("data_indep_permu_head_init_ortho_no_opt", True, 'ortho', True, False),
        #("data_dep_permu_head_init_ortho_no_opt", False, 'ortho', True, False),
        #("data_dep_permu_head_init_ortho_opt", 'ortho', True, True),
    ]
    
    for name, data_independent, 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 finetune_layer_which:
            activations_for_layer_a = activations_a[layer_idx]
            activations_for_layer_b = activations_b[layer_idx]

            result = align_attention_params_main(
                rng, params_a, aligned_params, layer_idx, num_heads, activations_for_layer_a, activations_for_layer_b, plot_path=None,
                init_method=init_method, permute_heads=permute_heads, optimize=optimize, method_name=name, data_independent=data_independent
            )

            # Update the aligned_params tree using the path returned by the alignment function.
            unfrozen_params = unfreeze(aligned_params)
            
            # Navigate to the parent dictionary of the MHA block
            temp_dict = unfrozen_params
            for key in result['mha_path'][:-1]:
                temp_dict = temp_dict[key]
            
            # Update the MHA block with the aligned parameters
            temp_dict[result['mha_path'][-1]] = result['aligned_params']
            
            aligned_params = freeze(unfrozen_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


        if optimize and plot_path:
            os.makedirs(plot_path, exist_ok=True)
            num_layers = len(finetune_layer_which)
            layers = finetune_layer_which
            metric_keys = ['objective_values', 'grad_norms', 'condition_nums']
            for metric_key in metric_keys:
                fig, axs = plt.subplots(num_layers + 1, 2, figsize=(20, 5 * (num_layers + 1)), sharex='col')
                for col in range(2):
                    if col == 0:
                        metrics_per_layer = layer_to_metrics_A
                        alignment_type = "Query/Key Alignment"
                    else:
                        metrics_per_layer = layer_to_metrics_B
                        alignment_type = "Value/Out Alignment"
                    
                    # Plot per-layer subplots
                    for row in range(num_layers):
                        layer = layers[row]
                        data_list = metrics_per_layer[layer][metric_key]
                        labels = [f"Head {i}" for i in range(num_heads)]
                        ax = axs[row, col]
                        for data, label in zip(data_list, labels):
                            if data:
                                final_val = data[-1]
                                ax.plot(data, label=f"{label} ({final_val:.4f})")
                        ax.set_title(f"Layer {layer}: {metric_key.replace('_', ' ').capitalize()} - {alignment_type}")
                        ax.set_xlabel('Iteration')
                        ax.set_ylabel(metric_key.replace('_', ' ').capitalize())
                        ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
                    
                    # Bottom row: mean across heads for all layers
                    ax = axs[num_layers, col]
                    data_list = []
                    labels = []
                    for layer in layers:
                        head_data = metrics_per_layer[layer][metric_key]
                        if head_data:
                            max_len = max(len(d) for d in head_data if d)
                            padded = []
                            for d in head_data:
                                if d:
                                    if len(d) < max_len:
                                        last = d[-1]
                                        padded.append(d + [last] * (max_len - len(d)))
                                    else:
                                        padded.append(d)
                            if padded:
                                mean_data = np.mean(padded, axis=0).tolist()
                                data_list.append(mean_data)
                                final_mean = mean_data[-1]
                                labels.append(f"Layer {layer} ({final_mean:.4f})")
                    for data, label in zip(data_list, labels):
                        ax.plot(data, label=label)
                    ax.set_title(f"All Layers Mean: {metric_key.replace('_', ' ').capitalize()} - {alignment_type}")
                    ax.set_xlabel('Iteration')
                    ax.set_ylabel(metric_key.replace('_', ' ').capitalize())
                    ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
                
                plt.tight_layout()
                save_path = os.path.join(plot_path, f"{name}_{metric_key}.png")
                plt.savefig(save_path, bbox_inches='tight')
                plt.close()

    return params_dict

#############RoPE#################
def get_rope_matrix(seq_len, d_head):
    """Generates RoPE rotation matrices R[m] of shape (seq_len, d_head/2, 2, 2)."""
    assert d_head % 2 == 0, "d_head must be even"
    # inv_freq: (d_head/2,)
    inv_freq = 1.0 / (10000 ** (jnp.arange(0, d_head, 2) / d_head))
    t = jnp.arange(seq_len)                 # (seq_len,)
    freqs = jnp.einsum('i,j->ij', t, inv_freq)   # (seq_len, d_head/2)

    cos_freqs = jnp.cos(freqs)   # (seq_len, h)
    sin_freqs = jnp.sin(freqs)   # (seq_len, h)

    # Build rotation matrices per position and subspace: shape (seq_len, h, 2, 2)
    # Each 2x2 is [[cos, -sin], [sin, cos]]
    R = jnp.stack(
        [
            jnp.stack([cos_freqs, -sin_freqs], axis=-1),  # (seq_len, h, 2) -> first row entries
            jnp.stack([sin_freqs,  cos_freqs], axis=-1),  # (seq_len, h, 2) -> second row entries
        ],
        axis=-2
    )  # After this stack: shape (seq_len, h, 2, 2)
    return R

@jax.jit
def apply_rope(x, R):
    """Applies RoPE to x of shape (B, L, D_k) using R shape (L, D_k/2, 2, 2)."""
    B, L, Dk = x.shape
    assert Dk % 2 == 0
    x_pairs = x.reshape((B, L, Dk//2, 2))   # (B, L, h, 2)
    # R must be (L, h, 2, 2)
    # einsum: 'b l h c, l h c r -> b l h r' -> back to (B, L, h, 2)
    x_rotated = jnp.einsum('blhc,lhcr->blhr', x_pairs, R)
    return x_rotated.reshape((B, L, Dk))

def compute_cost_matrix_presoftmax_rope(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_a, activations_b, alpha=0.5, epsilon=1e-8):
    B_a, L_a, D_a = activations_a.shape
    B_b, L_b, D_b = activations_b.shape
    assert B_a == B_b and L_a == L_b and D_a == D_b
    L, d_head = L_a, W_Q_a[0].shape[1]
    
    rope_matrices = get_rope_matrix(L, d_head)
    ones_col_a, ones_col_b = jnp.ones((B_a, L_a, 1)), jnp.ones((B_b, L_b, 1))
    X_tilde_a = jnp.concatenate([activations_a, ones_col_a], axis=-1)
    X_tilde_b = jnp.concatenate([activations_b, ones_col_b], axis=-1)
    sqrt_d = jnp.sqrt(float(d_head))
    C = np.zeros((num_heads, num_heads))

    # Pre-compute for model A
    S_bar_flat_a, V_flat_a = [], []
    for i in range(num_heads):
        tilde_W_Q_a_i, tilde_W_K_a_i, tilde_W_V_a_i = compute_extended_weights(W_Q_a[i], b_Q_a[i]), compute_extended_weights(W_K_a[i], b_K_a[i]), compute_extended_weights(W_V_a[i], b_V_a[i])
        Q_rope_a_i, K_rope_a_i = apply_rope(X_tilde_a @ tilde_W_Q_a_i, rope_matrices), apply_rope(X_tilde_a @ tilde_W_K_a_i, rope_matrices)
        S_a_i = jnp.einsum('bld,bmd->blm', Q_rope_a_i, K_rope_a_i) / sqrt_d
        S_bar_flat_a.append((S_a_i - jnp.mean(S_a_i, axis=2, keepdims=True)).flatten())
        V_flat_a.append((X_tilde_a @ tilde_W_V_a_i @ W_O_a[i]).flatten())

    # Pre-compute for model B
    S_bar_flat_b, V_flat_b = [], []
    for j in range(num_heads):
        tilde_W_Q_b_j, tilde_W_K_b_j, tilde_W_V_b_j = compute_extended_weights(W_Q_b[j], b_Q_b[j]), compute_extended_weights(W_K_b[j], b_K_b[j]), compute_extended_weights(W_V_b[j], b_V_b[j])
        Q_rope_b_j, K_rope_b_j = apply_rope(X_tilde_b @ tilde_W_Q_b_j, rope_matrices), apply_rope(X_tilde_b @ tilde_W_K_b_j, rope_matrices)
        S_b_j = jnp.einsum('bld,bmd->blm', Q_rope_b_j, K_rope_b_j) / sqrt_d
        S_bar_flat_b.append((S_b_j - jnp.mean(S_b_j, axis=2, keepdims=True)).flatten())
        V_flat_b.append((X_tilde_b @ tilde_W_V_b_j @ W_O_b[j]).flatten())
        
    for i in range(num_heads):
        for j in range(num_heads):
            dot_S = jnp.dot(S_bar_flat_a[i], S_bar_flat_b[j])
            norm_S_a, norm_S_b = jnp.linalg.norm(S_bar_flat_a[i]), jnp.linalg.norm(S_bar_flat_b[j])
            cost_S = 1.0 - (dot_S / (norm_S_a * norm_S_b + epsilon))
            dot_V = jnp.dot(V_flat_a[i], V_flat_b[j])
            norm_V_a, norm_V_b = jnp.linalg.norm(V_flat_a[i]), jnp.linalg.norm(V_flat_b[j])
            cost_V = 1.0 - (dot_V / (norm_V_a * norm_V_b + epsilon))
            C[i, j] = (alpha * cost_S + (1 - alpha) * cost_V)
    return C

def find_heads_permutation_rope(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_a, activations_b, alpha, data_independent):
    """Wrapper to find permutation using the RoPE cost matrix with model-specific activations."""
    if data_independent:
        C = compute_cost_matrix_data_independent(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, alpha=alpha)
    else:
        C = compute_cost_matrix_presoftmax_rope(
            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_a, activations_b, alpha
        )    
    row_ind, col_ind = linear_sum_assignment(to_numpy(C))
    print("RoPE Head Permutation:", {int(i): int(j) for i, j in zip(row_ind, col_ind)})
    return row_ind, col_ind

from scipy.optimize import minimize_scalar
from scipy.linalg import block_diag
from math import sqrt, cos, sin, atan2
import numpy as np

def solve_rope_qk_alignment(W_Q_a_i, b_Q_a_i, W_K_a_i, b_K_a_i,
                            W_Q_b_i, b_Q_b_i, W_K_b_i, b_K_b_i):
    """
    Solves for the G_RoPE alignment matrix U for a single head's QK weights.
    """
    tilde_W_Q_a = compute_extended_weights(W_Q_a_i, b_Q_a_i)
    tilde_W_K_a = compute_extended_weights(W_K_a_i, b_K_a_i)
    tilde_W_Q_b = compute_extended_weights(W_Q_b_i, b_Q_b_i)
    tilde_W_K_b = compute_extended_weights(W_K_b_i, b_K_b_i)
    
    D_k = tilde_W_Q_a.shape[1]
    assert D_k % 2 == 0, "Head dimension must be even for RoPE."

    U_blocks = []
    J = jnp.array([[0, -1], [1, 0]])

    for j in range(D_k // 2):
        # 1. Slice submatrices for the j-th 2D subspace
        sl = slice(2 * j, 2 * j + 2)
        Q_a_j, Q_b_j = tilde_W_Q_a[:, sl], tilde_W_Q_b[:, sl]
        K_a_j, K_b_j = tilde_W_K_a[:, sl], tilde_W_K_b[:, sl]

        # 2. Precompute constants
        N_Q = jnp.sum(Q_b_j**2)
        N_K = jnp.sum(K_b_j**2)
        C_Q = Q_a_j.T @ Q_b_j
        C_K = K_a_j.T @ K_b_j
        
        c_q = 0.5 * (jnp.trace(C_Q) + 1j * jnp.trace(C_Q @ J))
        c_k = 0.5 * (jnp.trace(C_K) + 1j * jnp.trace(C_K @ J))

        A = jnp.abs(c_q)**2
        B = jnp.abs(c_k)**2
        C = 2 * jnp.real(c_q * jnp.conj(c_k))

        # Convert all JAX/Device arrays to native Python/NumPy types for SciPy
        N_Q_f, N_K_f = float(N_Q), float(N_K)
        A_f, B_f, C_f = float(A), float(B), float(C)
        c_q_f = complex(c_q)
        c_k_f = complex(c_k)

        # Define the 1D scalar objective function using native floats
        def g_objective(x):
            x = float(x)
            # Protect the sqrt argument from tiny negative values due to roundoff
            inner_term = A_f * x + (B_f / x) + C_f
            safe_inner = max(inner_term, 1e-20)
            return x * N_Q_f + N_K_f / x - 4.0 * sqrt(safe_inner)

        # 4. Find the minimizer x* using robust bounds
        res = minimize_scalar(g_objective, bounds=(1e-8, 1e8), method='bounded')
        x_star = res.x

        # 5. Reconstruct the optimal 2x2 alignment matrix U_j using NumPy/math
        r_star = sqrt(x_star)
        combined_c = r_star * c_q_f + (1 / r_star) * c_k_f
        if abs(combined_c) < 1e-30:
            theta_star = 0.0
        else:
            theta_star = -atan2(combined_c.imag, combined_c.real)
        a = r_star * cos(theta_star)
        b = r_star * sin(theta_star)
        
        U_j = np.array([[a, -b], [b, a]])
        U_blocks.append(U_j)
        
    # 6. Assemble the full block-diagonal matrix U
    U_opt = block_diag(*U_blocks)
    condU = np.linalg.cond(U_opt)
    if condU > 1e12:
        # fallback: scale blocks to have minimum magnitude, or add small diag:
        eps = 1e-6
        U_opt = U_opt + eps * np.eye(U_opt.shape[0])
    return U_opt

def align_attention_params_main_rope(params_a, params_b, layer_idx, num_heads, 
                                     activations_for_layer_a, activations_for_layer_b, 
                                     init_method_vo='ortho', permu_heads=True, optimize_vo=True, alpha=0.5, data_independent=False):
    """
    Aligns a single MHA layer with RoPE using model-specific activations.
    COMPATIBLE with multiple model structures.
    """
    # Use the compatible extractor to get params and the update path ---
    params_a_extracted, _ = extract_attention_params(params_a, layer_idx)
    params_b_extracted, mha_path_b = extract_attention_params(params_b, layer_idx)
    params_a_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 'value', 'value_bias', 'out', 'out_bias'], params_a_extracted)}
    params_b_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 'value', 'value_bias', 'out', 'out_bias'], params_b_extracted)}
    W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_to_per_head(params_a_np, num_heads)
    W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_to_per_head(params_b_np, num_heads)

    if permu_heads:
        # --- Stage 1: Head Permutation (RoPE version) ---
        row_ind, col_ind = find_heads_permutation_rope(
            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_for_layer_a, activations_for_layer_b, alpha, data_independent
        )
        # --- Reorder heads of model B --- (No changes here)
        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]

    # --- Stage 2: Per-Head Parameter Alignment --- (No changes here)
    aligned_params_list = []
    for i in range(num_heads):
        # ... (rest of the function is identical)
        # QK Alignment (RoPE specific)
        U = solve_rope_qk_alignment(
            W_Q_a[i], b_Q_a[i], W_K_a[i], b_K_a[i],
            W_Q_b[i], b_Q_b[i], W_K_b[i], b_K_b[i]
        )
        U_inv = np.linalg.inv(U)
        
        W_Q_aligned = W_Q_b[i] @ U.T
        b_Q_aligned = b_Q_b[i] @ U.T
        W_K_aligned = W_K_b[i] @ U_inv
        b_K_aligned = b_K_b[i] @ U_inv
        
        # VO Alignment (Standard MHA logic, as it's unaffected by RoPE)
        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_V_b_i = compute_extended_weights(W_V_b[i], b_V_b[i])
        Y_O_b_i = W_O_b[i].T

        if init_method_vo == 'ortho':
            B_init = solve_orthogonal(Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i)
        else:
             B_init = np.identity(W_O_b[i].shape[0])

        if optimize_vo:
            B, _, _, _ = optimize_alignment(B_init, Y_O_a_i, Y_O_b_i, tilde_W_V_a_i, tilde_W_V_b_i)
        else:
            B = B_init

        B_inv = np.linalg.inv(B)
        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_list.append({
            '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}
        })

    # --- Reassemble Parameters --- (No changes here)
    query_kernel = np.stack([p['query']['kernel'] for p in aligned_params_list], axis=1)
    query_bias = np.stack([p['query']['bias'] for p in aligned_params_list], axis=0)
    key_kernel = np.stack([p['key']['kernel'] for p in aligned_params_list], axis=1)
    key_bias = np.stack([p['key']['bias'] for p in aligned_params_list], axis=0)
    value_kernel = np.stack([p['value']['kernel'] for p in aligned_params_list], axis=1)
    value_bias = np.stack([p['value']['bias'] for p in aligned_params_list], axis=0)
    out_kernel = np.stack([p['out']['kernel'] for p in aligned_params_list], axis=0)

    return {
        'aligned_params': {
            'query': {'kernel': jnp.array(query_kernel), 'bias': jnp.array(query_bias)},
            'key': {'kernel': jnp.array(key_kernel), 'bias': jnp.array(key_bias)},
            'value': {'kernel': jnp.array(value_kernel), 'bias': jnp.array(value_bias)},
            'out': {'kernel': jnp.array(out_kernel), 'bias': jnp.array(params_b_np['out_bias'])}
        },
        'mha_path': mha_path_b
    }

def matching_attn_rope(params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads, alpha=0.5):
    """
    Main function to align RoPE-based MHA layers using separate activations for each model.
    COMPATIBLE with multiple model structures.
    
    Args:
        params_a (dict): Parameters of the reference model A.
        params_b (dict): Parameters of the model B to be aligned.
        activations (dict): Dictionary mapping layer_idx to input activations.
        finetune_layer_which (list): List of layer indices to align.
        num_heads (int): Number of attention heads.
        alpha (float): Weighting factor for permutation cost matrix calculation.

    Returns:
        dict: A dictionary where keys are method names and values are the
              aligned parameters for model B.
    """
    params_dict = {}
    
    # Define a single configuration for RoPE alignment.
    # The VO part can optionally be optimized after orthogonal initialization.
    configurations = [
        ("data_indep_permu_head_init_ortho_no_opt", True, 'ortho', True, False),
        #("data_dep_permu_head_init_ortho_no_opt", False, 'ortho', True, False),
        #("data_dep_permu_head_init_ortho_opt", 'ortho', True, True),
    ]
    
    for name, data_independent, init_method_vo, permu_heads, optimize_vo in configurations:
        print(f"--- Running RoPE Alignment Configuration: {name} ---")
        
        # Start with a deepcopy of params_b to modify iteratively ---
        aligned_params = copy.deepcopy(params_b)

        for layer_idx in finetune_layer_which:
            print(f"Aligning Layer {layer_idx}...")
            
            activations_for_layer_a = activations_a[layer_idx]
            activations_for_layer_b = activations_b[layer_idx]

            # Call the main alignment function for a single RoPE MHA layer
            result = align_attention_params_main_rope(
                params_a=params_a,
                params_b=aligned_params,
                layer_idx=layer_idx,
                num_heads=num_heads,
                activations_for_layer_a=activations_for_layer_a,
                activations_for_layer_b=activations_for_layer_b,
                init_method_vo=init_method_vo,
                permu_heads=permu_heads,
                optimize_vo=optimize_vo,
                alpha=alpha,
                data_independent=data_independent
            )
            
            # Update the aligned_params tree using the path ---
            unfrozen_params = unfreeze(aligned_params)
            
            # Navigate to the parent dictionary of the MHA block
            temp_dict = unfrozen_params
            for key in result['mha_path'][:-1]:
                temp_dict = temp_dict[key]
            
            # Update the MHA block with the aligned parameters
            temp_dict[result['mha_path'][-1]] = result['aligned_params']
            
            aligned_params = freeze(unfrozen_params)


        # Sanity check
        total_sum = tree_util.tree_reduce(lambda acc, x: acc + jnp.sum(x), aligned_params, initializer=0)
        print(f"Finished configuration '{name}'. Total parameter sum: {total_sum:.4f}\n")
        
        params_dict[name] = aligned_params

    return params_dict


# this function performs matching on all possible heads permutations
# thus only support len(finetune_layer_which)==1 i.e. at 1 layer only
import itertools

layer_key_prefix = 'TransformerEncoderLayer'
attention_key = 'MultiHeadDotProductAttention_0'

def matching_attn_all_heads_permu(rng, params_a, params_b, finetune_layer_which, num_heads, plot_path, rope_use=False, activations_a=None, activations_b=None):
    """
    Performs matching for all possible head permutations for a single layer and records the
    optimal permutation for several data-dependent and independent methods.

    Args:
        rng: JAX random key.
        params_a: Parameters of the first model.
        params_b: Parameters of the second model.
        finetune_layer_which: A list containing the index of the layer to finetune (must have length 1).
        num_heads: The number of attention heads.
        plot_path: Path for saving plots (not used in this version but kept for consistency).
        rope_use: Boolean indicating if RoPE is used in the model.
        activations_a: A dictionary of activations from model A, keyed by layer index.
        activations_b: A dictionary of activations from model B, keyed by layer index.

    Returns:
        A tuple containing:
        - params_dict: Dictionary of aligned parameters for different settings and permutations.
        - heads_objective_values: Dictionary of objective values for each permutation.
        - heads_permutation_sol: Dictionary storing the optimal permutation for each calculation method.
    """
    assert len(finetune_layer_which) == 1, "This function only supports one layer at a time."

    params_dict = defaultdict(lambda: defaultdict(dict))
    heads_objective_values = defaultdict(lambda: defaultdict(dict))
    heads_permutation_sol = {}
    C_dict = {}

    layer_idx = finetune_layer_which[0]
    layer_key = f'{layer_key_prefix}_{layer_idx}'

    # --- Extract and reshape weights ---
    params_a_extracted = extract_attention_params(params_a, layer_key, attention_key)
    params_b_extracted = extract_attention_params(params_b, layer_key, attention_key)
    params_a_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 'value', 'value_bias', 'out', 'out_bias'], params_a_extracted)}
    params_b_np = {k: np.array(v) for k, v in zip(['key', 'key_bias', 'query', 'query_bias', 'value', 'value_bias', 'out', 'out_bias'], params_b_extracted)}
    W_Q_a, b_Q_a, W_K_a, b_K_a, W_V_a, b_V_a, W_O_a = reshape_to_per_head(params_a_np, num_heads)
    W_Q_b, b_Q_b, W_K_b, b_K_b, W_V_b, b_V_b, W_O_b = reshape_to_per_head(params_b_np, num_heads)

    # --- Define permutation settings ---
    head_permu_settings = [
        ("data-independent", True, None),
        #("data-dependent_acti-b-use", False, True),
        #("data-dependent_acti-b-notuse", False, False),
    ]

    alpha = 0.5  # Using a fixed alpha as in the original script

    # --- Calculate cost matrices and find optimal permutations for each method ---
    print("Calculating cost matrices and optimal permutations for each method...")
    for name, data_independent, activation_b_use in head_permu_settings:
        if data_independent:
            C = compute_cost_matrix_data_independent(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, alpha=alpha)
        else:
            # For data-dependent methods
            assert activations_a is not None, "Activations for model A must be provided for data-dependent methods."
            # Use model B's activations if specified, otherwise use model A's activations for both
            current_activations_b = activations_b[layer_idx] if activation_b_use and activations_b else activations_a[layer_idx]

            if rope_use:
                C = compute_cost_matrix_presoftmax_rope(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_a[layer_idx], current_activations_b, alpha=alpha)
            else:
                C = compute_cost_matrix_presoftmax(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_a[layer_idx], current_activations_b, alpha=alpha)

        C_dict[name] = C
        row_ind, col_ind = linear_sum_assignment(C)
        # Store the permutation mapping row_ind (model A) to col_ind (model B)
        # We sort by row_ind to ensure the permutation is always in the order [perm_for_head_0, perm_for_head_1, ...]
        permutation_solution = [int(col_ind[i]) for i in np.argsort(row_ind)]
        heads_permutation_sol[name] = ([int(_) for _ in row_ind], [int(_) for _ in  col_ind])
        print(f"  - Method '{name}': Optimal permutation is {permutation_solution}")


    # --- Iterate through all possible permutations for alignment and interpolation ---
    permutations_to_evaluate = list(itertools.permutations(range(num_heads)))

    # Collect unique permutations from heads_permutation_sol
    S = set()
    for name, _, _ in head_permu_settings:
        _, col_ind = heads_permutation_sol[name]
        col_ind = [int(_) for _ in col_ind] 
        S.add(tuple(col_ind))  # Convert to tuple for set compatibility

    # Initialize permutations to use
    permutations_to_use = set(S)
    max_permutations = 24

    # Sample additional permutations up to max_permutations
    while len(permutations_to_use) < max_permutations:
        perm = tuple(np.random.permutation(num_heads))
        permutations_to_use.add(perm)  # Set ensures no duplicates

    permutations_to_evaluate = permutations_to_use
    print(f"\nEvaluating all {len(permutations_to_evaluate)} possible head permutations...")

    # Define alignment settings
    alignment_settings = [
        ("init_ortho_no_opt", 'ortho', False),
        #("init_ortho_opt", 'ortho', True),
    ]

    for setting_name, init_method, optimize in alignment_settings:
        for perm_tuple in permutations_to_evaluate:
            perm = list(perm_tuple)
            # Create a fresh copy of model B parameters for this specific permutation
            params_b_permuted = copy.deepcopy(unfreeze(params_b))
            
            # Manually permute the attention heads of model B for the specified layer
            # This is done by re-indexing the weight tensors according to `perm`
            params_b_layer = params_b_permuted[layer_key][attention_key]
            perm_array = np.array(perm)
            
            # Permute dimensions related to heads
            params_b_layer['query']['kernel'] = params_b_layer['query']['kernel'][:, perm_array, :]
            params_b_layer['query']['bias'] = params_b_layer['query']['bias'][perm_array, :]
            params_b_layer['key']['kernel'] = params_b_layer['key']['kernel'][:, perm_array, :]
            params_b_layer['key']['bias'] = params_b_layer['key']['bias'][perm_array, :]
            params_b_layer['value']['kernel'] = params_b_layer['value']['kernel'][:, perm_array, :]
            params_b_layer['value']['bias'] = params_b_layer['value']['bias'][perm_array, :]
            params_b_layer['out']['kernel'] = params_b_layer['out']['kernel'][perm_array, :, :]

            # Align the permuted model B to model A
            if rope_use:
                aligned_result = align_attention_params_main_rope(
                    params_a, params_b_permuted, layer_idx, num_heads,
                    activations_a[layer_idx],
                    activations_b[layer_idx],
                    init_method_vo=init_method, permu_heads=False, optimize_vo=optimize, alpha=alpha
                )
                # The rope main function returns the aligned MHA block directly
                result_to_store = aligned_result
            else:
                # Standard alignment. `permute_heads` is False because we do it manually above.
                result_dict = align_attention_params_main(
                    rng, params_a, params_b_permuted, layer_idx, num_heads,
                    activations_a[layer_idx],
                    activations_b[layer_idx],
                    plot_path=None, init_method=init_method,
                    permute_heads=False, optimize=optimize, alpha=alpha
                )
                result_to_store = result_dict

            perm_str = str([int(_) for _ in perm])
            params_dict[setting_name][perm_str] = result_to_store

            # --- Calculate objective cost for this permutation against each method's cost matrix ---
            for method_name, C_matrix in C_dict.items():
                # The cost is the sum of C[i, j] for the mapping i -> perm[i]
                total_cost = sum(C_matrix[i, perm[i]] for i in range(num_heads))
                heads_objective_values[method_name][setting_name][perm_str] = float(total_cost)

    print("Finished evaluating all permutations.")
    return params_dict, heads_objective_values, heads_permutation_sol

##########Transformers matching
import copy
import jax.numpy as jnp
import numpy as np
from flax.core import unfreeze, freeze
from scipy.optimize import linear_sum_assignment

# --- Helper Functions (from your draft and my previous code) ---

def get_nested_item(d, keys):
    """Accesses a nested dictionary item using a tuple of keys."""
    for key in keys:
        d = d[key]
    return d

def set_nested_item(d, keys, value):
    """Sets a value in a nested dictionary using a tuple of keys."""
    current = d
    for key in keys[:-1]:
        current = current[key]
    current[keys[-1]] = value

def extract_ffn_params(params, layer_idx):
    """
    Extracts FFN parameters from different, known model structures in a compatible way.
    This function detects the model type and constructs the correct path to the
    FFN parameters for a given layer index.
    Args:
        params: The Flax parameter tree.
        layer_idx: The integer index of the transformer layer.
    Returns:
        A tuple containing:
        - W1, b1, W2 (FFN weights and biases for Dense_0 and Dense_1).
        - dense0_path, dense1_path (tuples representing the nested paths to Dense_0 and Dense_1).
    """
    # Detect cifar_vit style: params['TransformerEncoderLayer_...']
    if f'TransformerEncoderLayer_{layer_idx}' in params:
        base_path_tuple = (f'TransformerEncoderLayer_{layer_idx}',)
        dense0_path = base_path_tuple + ('Dense_0',)
        dense1_path = base_path_tuple + ('Dense_1',)
    # Detect vit-jax style: params['Transformer']['encoderblock_...']
    elif 'Transformer' in params and f'encoderblock_{layer_idx}' in params['Transformer']:
        # Note: vit-jax often nests the MLP in its own block, e.g., 'MlpBlock_0'
        # We check for its existence for robustness.
        encoder_block = get_nested_item(params, ('Transformer', f'encoderblock_{layer_idx}'))
        mlp_key = next((k for k in encoder_block if 'MlpBlock' in k), None)
        if mlp_key:
             base_path_tuple = ('Transformer', f'encoderblock_{layer_idx}', mlp_key)
        else: # Fallback if no explicit MlpBlock
             base_path_tuple = ('Transformer', f'encoderblock_{layer_idx}')
        dense0_path = base_path_tuple + ('Dense_0',)
        dense1_path = base_path_tuple + ('Dense_1',)
    else:
        raise KeyError(f"Could not find a known path for FFN layer {layer_idx} in the provided params.")
    
    dense0_block = get_nested_item(params, dense0_path)
    dense1_block = get_nested_item(params, dense1_path)
    
    W1 = dense0_block['kernel']
    b1 = dense0_block['bias']
    W2 = dense1_block['kernel']
    
    return W1, b1, W2, dense0_path, dense1_path

# --- Core Functions ---

def matching_transformer_ffn(params_a, params_b, finetune_layer_which):
    """
    Aligns the FFN components of two models for the specified layers.

    This function computes the optimal permutation of hidden neurons in the FFN
    of model B to match model A. The permutation is found by solving a Linear
    Assignment Problem (LAP) where the cost is the sum of squared L2 distances
    between the incoming (weights + bias) and outgoing weights of each neuron pair.
    
    Args:
        params_a (dict): Parameters of the reference model A.
        params_b (dict): Parameters of the model B to be aligned.
        finetune_layer_which (list): List of layer indices to align.
    
    Returns:
        dict: The aligned parameters for model B.
    """
    aligned_params_b = copy.deepcopy(params_b)

    for layer_idx in finetune_layer_which:
        # Extract parameters and paths for both models
        W1_a, b1_a, W2_a, _, _ = extract_ffn_params(params_a, layer_idx)
        W1_b, b1_b, W2_b, dense0_path_b, dense1_path_b = extract_ffn_params(aligned_params_b, layer_idx)
        
        # Convert to NumPy for computation
        W1_a, b1_a, W2_a = np.array(W1_a), np.array(b1_a), np.array(W2_a)
        W1_b, b1_b, W2_b = np.array(W1_b), np.array(b1_b), np.array(W2_b)

        D_hidden = W1_a.shape[1]
        assert W1_b.shape[1] == D_hidden, f"FFN hidden dimensions for layer {layer_idx} must match."
        
        # Compute cost matrix C
        C = np.zeros((D_hidden, D_hidden), dtype=np.float32)
        for i in range(D_hidden):
            # Incoming weights and bias for neuron i of model A
            in_a = np.concatenate([W1_a[:, i], [b1_a[i]]])
            # Outgoing weights for neuron i of model A
            out_a = W2_a[i, :]
            for j in range(D_hidden):
                in_b = np.concatenate([W1_b[:, j], [b1_b[j]]])
                out_b = W2_b[j, :]
                
                # Cost is the sum of squared Euclidean distances
                cost = np.linalg.norm(in_a - in_b)**2 + np.linalg.norm(out_a - out_b)**2
                C[i, j] = cost
        
        # Solve LAP. `col_ind` gives the permutation for model B's neurons.
        row_ind, col_ind = linear_sum_assignment(C)
        
        # Permute the weights of model B according to the solution
        W1_aligned = W1_b[:, col_ind]
        b1_aligned = b1_b[col_ind]
        W2_aligned = W2_b[col_ind, :]
        
        # Update the parameter dictionary for model B
        unfrozen_params = unfreeze(aligned_params_b)
        set_nested_item(unfrozen_params, dense0_path_b + ('kernel',), jnp.array(W1_aligned))
        set_nested_item(unfrozen_params, dense0_path_b + ('bias',), jnp.array(b1_aligned))
        set_nested_item(unfrozen_params, dense1_path_b + ('kernel',), jnp.array(W2_aligned))
        aligned_params_b = freeze(unfrozen_params)
    
    return aligned_params_b

def matching_transformer_block(rng, params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads, rope_use=False, plot_path=None):
    """
    Aligns entire Transformer blocks by sequentially aligning their MHA and FFN components.

    This function first calls the appropriate MHA alignment function (`matching_attn` or
    `matching_attn_rope`), which returns a dictionary of aligned parameters for various
    configurations. It then iterates through this dictionary, applying the FFN
    alignment to each MHA-aligned model.
    
    Args:
        rng: JAX random key.
        params_a (dict): Parameters of the reference model A.
        params_b (dict): Parameters of the model B to be aligned.
        activations_a (dict): Dictionary of activations from model A, keyed by layer index.
        activations_b (dict): Dictionary of activations from model B, keyed by layer index.
        finetune_layer_which (list): List of layer indices to align.
        num_heads (int): Number of attention heads.
        rope_use (bool): If True, use RoPE-specific MHA alignment.
        plot_path (str, optional): Path for saving diagnostic plots.
    
    Returns:
        dict: A dictionary where keys are configuration names (e.g., 'data_indep_...`) 
              and values are the fully aligned (MHA + FFN) parameter dictionaries.
    """
    print("--- Starting Transformer Block Alignment ---")
    
    # Step 1: Align the MHA component for all specified layers and configurations.
    print("\nStep 1: Aligning Multi-Head Attention components...")
    if rope_use:
        mha_aligned_params_dict = matching_attn_rope(params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads)
    else:
        mha_aligned_params_dict = matching_attn(rng, params_a, params_b, activations_a, activations_b, finetune_layer_which, num_heads, plot_path)
    print("MHA alignment complete.")

    # Step 2: For each MHA-aligned model, align its FFN component.
    print("\nStep 2: Aligning Feed-Forward Network components for each configuration...")
    fully_aligned_params_dict = {}
    for config_name, mha_aligned_params in mha_aligned_params_dict.items():
        print(f"  - Aligning FFN for configuration: '{config_name}'")
        fully_aligned_params = matching_transformer_ffn(params_a, mha_aligned_params, finetune_layer_which)
        fully_aligned_params_dict[config_name] = fully_aligned_params
    
    print("FFN alignment complete.")
    print("\n--- Transformer Block Alignment Finished ---")
    
    return fully_aligned_params_dict