Project: 041_Unlocking_the_Power_of_Function_Vectors_for_Charac | Function: compute_function_vector (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 function vectors by aggregating attention head activations weighted by their indirect effects

**KEY TERMS:** {
  "mean_activations": "h̄Tlk",
  "indirect_effect": "causal effect score",
  "top_heads": "attention heads set S",
  "function_vector": "θT"
}

**CONTEXT CODE:**
    rand_tokens_repeat =  torch.concat((rand_tokens, rand_tokens), dim=1).to(model.device)

    output = model(rand_tokens_repeat, output_attentions = True)
    attentions = torch.vstack(output.attentions) # (n_layers, n_heads, tokens, tokens)

    score = attentions.diagonal(1-seq_len, dim1=-2, dim2=-1)
    score_per_head = score.reshape(model_config['n_layers'],batch_size,model_config['n_heads'],-1).mean(axis=-1).mean(axis=1).cpu().T
    
    return score_per_head

def compute_function_vector(mean_activations, indirect_effect, model, model_config, n_top_heads = 10, token_class_idx=-1, gen=False):
    """
        Computes a "function vector" vector that communicates the task observed in ICL examples used for downstream intervention.
        
        Parameters:
        mean_activations: tensor of size (Layers, Heads, Tokens, head_dim) containing the average activation of each head for a particular task
        indirect_effect: tensor of size (N, Layers, Heads, class(optional)) containing the indirect_effect of each head across N trials
        model: huggingface model being used
        model_config: contains model config information (n layers, n heads, etc.)
        n_top_heads: The number of heads to use when computing the summed function vector
        token_class_idx: int indicating which token class to use, -1 is default for last token computations

        Returns:
        function_vector: vector representing the communication of a particular task
        top_heads: list of the top influential heads represented as tuples [(L,H,S), ...], (L=Layer, H=Head, S=Avg. Indirect Effect Score)         
    """
    raise NotImplementedError("Function implementation has been masked for testing")

def compute_universal_function_vector(mean_activations, model, model_config, n_top_heads=10):
    """
        Computes a "function vector" vector that communicates the task observed in ICL examples used for downstream intervention
        using the set of heads with universally highest causal effect computed across a set of ICL tasks
        
        Parameters:
        mean_activations: tensor of size (Layers, Heads, Tokens, head_dim) containing the average activation of each head for a particular task
        model: huggingface model being used
        model_config: contains model config information (n layers, n heads, etc.)
        n_top_heads: The number of heads to use when computing the function vector

        Returns:
        function_vector: vector representing the communication of a particular task
        top_heads: list of the top influential heads represented as tuples [(L,H,S), ...], (L=Layer, H=Head, S=Avg. Indirect Effect Score)         
    """
    model_resid_dim = model_config['resid_dim']
    model_n_heads = model_config['n_heads']
    model_head_dim = model_resid_dim//model_n_heads
    device = model.device

    # Universal Set of Heads
    
    if 'gpt2' in model_config['name_or_path']:
        top_heads =[
            (21, 13, 0.0802), (25, 13, 0.0368), (20, 5, 0.035), (21, 19, 0.0346), (19, 18, 0.0321), (29, 2, 0.0319), 
            (20, 16, 0.0317), (23, 16, 0.03), (29, 0, 0.0285), (22, 15, 0.0264), (26, 16, 0.0246), (23, 20, 0.0234), 
            (20, 19, 0.0215), (23, 11, 0.0206), (23, 5, 0.0197), (24, 11, 0.0191), (17, 16, 0.0189), (16, 21, 0.0178), 
            (27, 15, 0.0163), (32, 20, 0.0161), (17, 1, 0.0152), (28, 2, 0.0151), (13, 0, 0.0148), (19, 20, 0.0137), 
            (19, 10, 0.0127), (22, 11, 0.0126), (26, 22, 0.0122), (25, 19, 0.0118), (36, 8, 0.0112), (28, 19, 0.0111), 
            (24, 9, 0.0104), (39, 18, 0.0099), (19, 5, 0.0094), (18, 22, 0.0092), (33, 9, 0.009), (21, 3, 0.0085), 
            (34, 19, 0.0081), (26, 6, 0.0079), (23, 8, 0.0073), (17, 6, 0.0072), (24, 21, 0.0071), (42, 20, 0.007), 
            (21, 22, 0.007), (34, 22, 0.0069), (16, 3, 0.0068)
        ]
    elif 'gpt-j' in model_config['name_or_path']:
        top_heads = [(15, 5, 0.0587), (9, 14, 0.0584), (12, 10, 0.0526), (8, 1, 0.0445), (11, 0, 0.0445), (13, 13, 0.019), (8, 0, 0.0184), (14, 9, 0.016), (9, 2, 0.0127), (24, 6, 0.0113), (15, 11, 0.0092),
                     (6, 6, 0.0069), (14, 0, 0.0068), (17, 8, 0.0068), (21, 2, 0.0067), (10, 11, 0.0066), (11, 2, 0.0057), (17, 0, 0.0054), (20, 11, 0.0051), (23, 0, 0.0047), (20, 0, 0.0046), (15, 7, 0.0045),
                     (27, 2, 0.0045), (21, 15, 0.0044), (11, 4, 0.0044), (18, 6, 0.0043), (9, 6, 0.0042), (4, 12, 0.004), (11, 15, 0.004), (20, 2, 0.0036), (10, 0, 0.0035), (16, 9, 0.0031), (11, 14, 0.0031),
                     (12, 4, 0.003), (9, 7, 0.003), (18, 3, 0.003), (19, 5, 0.003), (22, 5, 0.0027), (25, 3, 0.0026), (18, 9, 0.0025)]
    elif 'Llama-2-7b' in model_config['name_or_path']:
        top_heads = [(14, 1, 0.0391), (11, 2, 0.0225), (9, 25, 0.02), (12, 15, 0.0196), (12, 28, 0.0191), (13, 7, 0.0171), (11, 18, 0.0152), (12, 18, 0.0113), (16, 10, 0.007), (14, 16, 0.007),
                     (14, 14, 0.0048), (16, 1, 0.0042), (18, 1, 0.0042), (19, 16, 0.0041), (13, 30, 0.0034), (18, 26, 0.0032), (14, 7, 0.0032), (16, 0, 0.0031), (16, 29, 0.003), (29, 30, 0.003),
                     (16, 6, 0.0029), (15, 11, 0.0027), (12, 11, 0.0026), (11, 22, 0.0023), (16, 19, 0.0021), (15, 23, 0.002), (16, 20, 0.0019), (15, 9, 0.0019), (17, 28, 0.0019), (14, 18, 0.0018),