Project: 061_REEF_Representation_Encoding_Fingerprints_for_Larg | Function: linear_CKA (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Computes the Centered Kernel Alignment (CKA) similarity between two sets of representations using linear kernel. This is the core REEF fingerprinting algorithm that measures similarity between suspect and victim model representations.

**KEY TERMS:** {
  "X": "suspect model representations",
  "Y": "victim model representations",
  "hsic": "HSIC similarity score",
  "var1": "sqrt(HSIC(X,X))",
  "var2": "sqrt(HSIC(Y,Y))"
}

**CONTEXT CODE:** import os
import argparse
try:
    from utils import seed_torch
except ImportError:
    def seed_torch(seed):
        import torch
        torch.manual_seed(seed)

try:
    from generate_activations import load_acts
except ImportError:
    # For testing without real data loading
    def load_acts(*args, **kwargs):
        import torch
        return torch.randn(10, 128)
import numpy as np
import math
import torch


class CKA(object):
    def __init__(self):
        pass 
    
    def centering(self, K):
        n = K.shape[0]
        unit = np.ones([n, n])
        I = np.eye(n)
        H = I - unit / n
        return np.dot(np.dot(H, K), H) 

    def rbf(self, X, sigma=None):
        GX = np.dot(X, X.T)
        KX = np.diag(GX) - GX + (np.diag(GX) - GX).T
        if sigma is None:
            mdist = np.median(KX[KX != 0])
            sigma = math.sqrt(mdist)
        KX *= - 0.5 / (sigma * sigma)
        KX = np.exp(KX)
        return KX
 
    def kernel_HSIC(self, X, Y, sigma):
        return np.sum(self.centering(self.rbf(X, sigma)) * self.centering(self.rbf(Y, sigma)))

    def linear_HSIC(self, X, Y):
        L_X = X @ X.T
        L_Y = Y @ Y.T
        return np.sum(self.centering(L_X) * self.centering(L_Y))

    def linear_CKA(self, X, Y):
        raise NotImplementedError("Function implementation has been masked for testing")
    
    def kernel_CKA(self, X, Y, sigma=None):
        hsic = self.kernel_HSIC(X, Y, sigma)
        var1 = np.sqrt(self.kernel_HSIC(X, X, sigma))
        var2 = np.sqrt(self.kernel_HSIC(Y, Y, sigma))

        return hsic / (var1 * var2)

    
class CudaCKA(object):
    def __init__(self, device):
        self.device = device
    
    def centering(self, K):
        n = K.shape[0]
        unit = torch.ones([n, n], device=self.device)
        I = torch.eye(n, device=self.device)
        H = I - unit / n
        return torch.matmul(torch.matmul(H, K), H)  

    def rbf(self, X, sigma=None):
        GX = torch.matmul(X, X.T)
        KX = torch.diag(GX) - GX + (torch.diag(GX) - GX).T
        if sigma is None:
            mdist = torch.median(KX[KX != 0])
            sigma = math.sqrt(mdist)
        KX *= - 0.5 / (sigma * sigma)
        KX = torch.exp(KX)
        return KX

    def kernel_HSIC(self, X, Y, sigma):
        return torch.sum(self.centering(self.rbf(X, sigma)) * self.centering(self.rbf(Y, sigma)))

    def linear_HSIC(self, X, Y):
        L_X = torch.matmul(X, X.T)
        L_Y = torch.matmul(Y, Y.T)
        return torch.sum(self.centering(L_X) * self.centering(L_Y))

    def linear_CKA(self, X, Y):
        hsic = self.linear_HSIC(X, Y)
        var1 = torch.sqrt(self.linear_HSIC(X, X))
        var2 = torch.sqrt(self.linear_HSIC(Y, Y))

        return hsic / (var1 * var2)

    def kernel_CKA(self, X, Y, sigma=None):
        hsic = self.kernel_HSIC(X, Y, sigma)
        var1 = torch.sqrt(self.kernel_HSIC(X, X, sigma))
        var2 = torch.sqrt(self.kernel_HSIC(Y, Y, sigma))
        return hsic / (var1 * var2)


if __name__ == "__main__":
    """
    Compute the CKA of the base model (victim model) and the test model (suspect model)
    """
    parser = argparse.ArgumentParser(description="Generate activations for statements in a dataset")
    parser.add_argument("--base_model", default="llama-2-7b")
    parser.add_argument("--base_layers", nargs='+', type=int, default=[18], help="Layers to load activations; can also set 24/32/40/80 to load all layers' activations according to the model size")
    parser.add_argument("--test_model", default="vicuna-7b-v1.5")
    parser.add_argument("--test_layers", nargs='+', type=int, default=[18], help="Layers to load activations; can also set 24/32/40/80 to load all layers' activations according to the model size")

    parser.add_argument("--datasets", nargs='+', default=['truthfulqa-200'], help="Names of datasets, without .csv extension")
    parser.add_argument("--output_dir", default="activations", help="Directory to save activations to")
    parser.add_argument("--device", default="cuda")

    parser.add_argument('--seed', default=2024, type=int)
    parser.add_argument('--permutation', action='store_true')
    parser.add_argument('--scaling', action='store_true')
    parser.add_argument('--l2', action='store_true')  
    parser.add_argument('--noise', action='store_true') 

    parser.add_argument("--load_in_8bit", action='store_true')
    parser.add_argument("--load_in_4bit", action='store_true')
    args = parser.parse_args()


    device = args.device
    seed_torch(args.seed)
    dataset_name = f'{args.datasets[0]}'

    if args.base_layers[0] in [24,32,40,80]:
        base_layers = list(range(args.base_layers[0]))
    else:
        base_layers = [int(layer) for layer in args.base_layers]

    test_layers = list(range(args.test_layers[0]))

    base_layer_acts = []
    for layer in base_layers:    
        acts = load_acts(dataset_name, args.base_model, layer=layer, center=True, scale=True, device=args.device, acts_dir=args.output_dir)
        base_layer_acts.append(acts)

    base_layer_acts = torch.stack(base_layer_acts).to(args.device) # [layer_num, data_num, feature_dim]


    test_model = args.test_model
    if args.load_in_8bit:


Complete the function implementation replacing "raise NotImplementedError".