Project: 019_Do_I_Know_This_Entity_Knowledge_Awareness_and_Hall | Function: steer_sae_latents (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:** SAE latent activation steering implementation for causal intervention testing using directional activation addition to validate entity recognition directions

**KEY TERMS:** {
  "activation": "x input activation tensor",
  "direction": "dⱼ SAE decoder latent direction",
  "coeff_value": "α steering coefficient",
  "pos": "token positions for intervention",
  "unsqueeze": "tensor reshaping for broadcasting"
}

**CONTEXT CODE:** import torch
import einops
import gc
import copy
import plotly.graph_objects as go
from dataset.load_data import load_wikidata_queries
from collections import defaultdict
import numpy as np
from typing import List, Union, Optional, Literal, Tuple
from torch import Tensor
import torch.nn.functional as F
from fancy_einsum import einsum
from jaxtyping import Float
from transformer_lens import HookedTransformer, HookedTransformerConfig, FactoredMatrix, ActivationCache
import transformer_lens.patching as patching
from functools import partial
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
from adjustText import adjust_text
from tqdm import tqdm
import json

import random
random_seed = 42
random.seed(random_seed)

### Hooks
def steer_sae_latents(activation, hook, direction, pos, coeff_value=1):
        
    raise NotImplementedError("Function implementation has been masked for testing")
def ablate_sae_latents(activation, hook, direction, pos):
        # Project away a direction from the residual stream
        if activation.shape[1]==1:
            # generating
            return activation

        if pos != 'all':
            if isinstance(pos[0], list):
                for batch_idx, p in enumerate(pos):
                    activation[batch_idx, p, :] -= (activation[batch_idx, p, :] @ direction.unsqueeze(-1)) * direction
            else:
                activation[:, pos, :] -= (activation[:, pos, :] @ direction.unsqueeze(-1)) * direction
        else:
            # activation -> batch, seq_len, d_model
            # direction -> d_model
            activation[:, :, :] -= (activation[:, :, :] @ direction.unsqueeze(-1)) * direction

        return activation


def generation_steered_latents(model, ids, pos, steering_latents=None, ablate_latents=None, coeff_value=1, max_new_tokens=30):
    if steering_latents is not None:
        # TODO: change this
        steer_fwd_hooks = [(f"blocks.{layer}.hook_resid_pre", partial(steer_sae_latents,
                                                pos=pos,
                                                direction=direction,
                                                coeff_value=coeff_value))
                            for layer, latent_idx, mean_act, direction in steering_latents]
    else:
        steer_fwd_hooks = []

    if ablate_latents is not None:
        ablate_fwd_hooks = [(f"blocks.{layer}.hook_resid_pre", partial(ablate_sae_latents,
                                                pos=pos,
                                                direction=direction))
                            for layer, latent_idx, mean_act, direction in ablate_latents]
    else:
        ablate_fwd_hooks = []
    
    for hook_filter, hook_fn in steer_fwd_hooks:
        model.add_hook(hook_filter, hook_fn, "fwd")

    for hook_filter, hook_fn in ablate_fwd_hooks:
        model.add_hook(hook_filter, hook_fn, "fwd")
    generations = model.generate(ids, max_new_tokens=max_new_tokens, do_sample=False)
    steered_generations = [model.to_string(generation) for generation in generations]
    return steered_generations

def cache_steered_latents(model, ids, pos, steering_latents: List[Tuple[int, float, Tensor]]=None, ablate_latents: List[Tuple[int, float, Tensor]]=None, coeff_value=1):
        
    if steering_latents is not None:
        steer_fwd_hooks = [(f"blocks.{layer}.hook_resid_pre", partial(steer_sae_latents,
                                                pos=pos,
                                                direction=direction,
                                                coeff_value=mean_act if coeff_value == 'mean' else coeff_value))
                            for layer, latent_idx, mean_act, direction in steering_latents]
    else:
        steer_fwd_hooks = []

    if ablate_latents is not None:
        ablate_fwd_hooks = [(f"blocks.{layer}.hook_resid_pre", partial(ablate_sae_latents,
                                                pos=pos,
                                                direction=direction))
                            for layer, latent_idx, mean_act, direction in ablate_latents]
    else:
        ablate_fwd_hooks = []
    
    for hook_filter, hook_fn in steer_fwd_hooks:
        model.add_hook(hook_filter, hook_fn, "fwd")

    for hook_filter, hook_fn in ablate_fwd_hooks:
        model.add_hook(hook_filter, hook_fn, "fwd")
    output = model.run_with_cache(ids, return_type="logits")
    return output



######
def get_batch_pos(batch_entity_pos, pos_type, tokenized_prompts):
    """
    Computes returns the positions to se for steering in the tokenized prompts.

    Args:
        batch_entity_pos (List[List[int]]): A list of lists where each sublist contains the positions of entity tokens in a prompt.
        pos_type (str): The type of position to compute. It can be one of the following:
            - 'all': Use all positions.
            - 'entity': Use the positions of the entity tokens.
            - 'entity_and_eoi': Use the positions of the entity tokens and the next 6 tokens.
            - 'entity_to_end': Use the positions of the entity tokens and all tokens to the end of the sequence.
            - 'entity_last': Use only the last position of the entity tokens.
            - 'entity_last_to_end': Use the last position of the entity tokens and all tokens to the end of the sequence.
        tokenized_prompts (List[List[int]]): A list of lists where each sublist contains the tokenized prompt.

    Returns:
        List[List[int]]: A list of lists where each sublist contains the computed positions.
    """



Complete the function implementation replacing "raise NotImplementedError".