extra_materials/graph/Feature_Graph_Blip.py

import torch as t
from torch import Tensor
from .Graph_Template import Graph, GraphName, Node, Index
from typing import List, Tuple, Dict, Union, Callable, Any
from tqdm import tqdm
from transformer_lens.hook_points import HookPoint
from transformer_lens import (
    utils,
    HookedTransformer,
    ActivationCache,
)
from copy import deepcopy
from collections import OrderedDict, defaultdict
from warnings import warn
from contextlib import contextmanager
from model.hooked_blip import HookedSAEBlipConditionalGeneration
from functools import partial
from itertools import product
from collections import OrderedDict
from .Graph_utils import nested_dict_to_string
from .Sparse_act import SparseAct
from torch import sparse_coo_tensor
import einops
from .vision_sae_wrapper import *

def cache_to_sparseact(
    cache: Dict[str, Tensor] | ActivationCache, 
    act_name: str, 
    res_name: str | None = None, 
    resc_name: str | None = None,
) -> SparseAct:
    return SparseAct(
        act=cache[act_name],
        res=cache[res_name] if res_name is not None else None,
        resc=cache[resc_name] if resc_name is not None else None,
    )
    
def interpolate(start: Tensor, end: Tensor, frac: float,) -> Tensor:
    assert 0 <= frac <= 1, "frac must be in [0, 1]"
    return start * frac + end * (1 - frac)

def create_list(dict: Dict[str, Tensor], name: str) -> Dict[str, Tensor]:
    if name not in dict:
        dict[name] = 0 # type: ignore
    return dict
    
class ConnectionNode(Node):
    def __init__(self, name: str):
        self._name = name
        
    @property
    def name(self) -> str:
        return self._name
    
    def __eq__(self, other) -> bool:
        if isinstance(other, ConnectionNode):
            return self._name == other.name
        elif isinstance(other, str):
            return self._name == other
        else:
            raise NotImplementedError("other is not an instance of ConnectionNode or str")
    
    def __repr__(self) -> str:
        return self.name
    
    def __hash__(self) -> int:
        return hash(self.name)
    
class ConnectionIndex(Index):
    def __init__(
        self,
        list_index: tuple[int|None, ...] | None = None,
    ):
        if list_index is None:
            self.list_index = (None,)
        else:
            for index in list_index:
                assert type(index) == int or index == None, "index is not an instance of int or None"
            self.list_index = list_index
        
    @property
    def as_index(self) -> Tuple[int | slice, ...]:
        return tuple(slice(None) if x is None else x for x in self.list_index) # for indexing
    
    def __repr__(self) -> str:
        ret = "["
        for idx, x in enumerate(self.list_index):
            if idx > 0:
                ret += ", "
            if x is None:
                ret += ":"
            elif type(x) == int:
                ret += str(x)
            else:
                raise NotImplementedError(x)
        ret += "]"
        return ret
    
    def __eq__(self, other) -> bool:
        if isinstance(other, ConnectionIndex):
            return self.list_index == other.list_index
        elif isinstance(other, tuple):
            return self.list_index == other
        else:
            raise NotImplementedError("other is not an instance of ConnectionIndex or tuple")
    
    def __hash__(self) -> int:
        return hash(self.list_index)
    
class FeatureIndex(ConnectionIndex):
    def __init__(
        self, 
        idx: Tuple[int, ...] | List[int], 
        length = 2,
    ):
        self.idx = tuple(idx)
        super().__init__(return_idx(idx, length))
        
class FeatureErrorIndex(ConnectionIndex):
    def __init__(
        self, 
        idx: Tuple[int, ...] | List[int], 
        length = 2,
    ):
        self.idx = tuple(idx)
        super().__init__(return_idx(idx, length))
    
class ErrorIndex(ConnectionIndex):
    def __init__(
        self, 
        idx: Tuple[int, ...] | List[int],
        length = 1,
    ):
        self.idx = tuple(idx)
        super().__init__(return_idx(idx, length))

def sae_hook_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_acts_post' # e.g. ...hook_resid_pre.hook_sae_acts_post

def error_term_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_error'

def output_hook_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_output'

def input_hook_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_input'

def recons_hook_name(sae_name: str) -> str:
    return f'{sae_name}.hook_sae_recons'

def revert_hook_name(sae_name: str) -> str:
    return ".".join(sae_name.split(".")[:-1]) # e.g. ...hook_resid_pre.hook_sae_acts_post -> ...hook_resid_pre

def return_idx(idx: Tuple | List, length: int = 2) -> Tuple:
    assert len(idx) <= length, "The length of idx must be less than or equal to length"
    return tuple([None for _ in range(length - len(idx))] + list(idx))
    

class Feature_Graph_Blip(Graph):
    def __init__(
        self,
        model: HookedSAEBlipConditionalGeneration,
        text_saes: Dict[int, List[Tuple[str, Any]]], # {layer: list[{hook_position: HookedSAE}]}, can define granularity here
        vision_saes: Dict[int, List[Tuple[str, Any]]],
        use_error_term: bool = False,
    ):
        '''
        '''
        self.model = model
        self.cfg = model.cfg
        self.use_error_term = use_error_term
        self.device = self.cfg.device
        
        self.n_layers = self.cfg.n_layers
        # self.n_heads = self.cfg.n_heads
        
        assert len(text_saes) == self.n_layers, "please provide SAEs"
        self.text_saes = text_saes
        self.vision_saes = vision_saes
        self.dict_vision_saes = self.extract_saes(vision_saes)
        self.dict_text_saes = self.extract_saes(text_saes)
        self.dict_saes: Dict[str, Any] = self.dict_vision_saes | self.dict_text_saes
        
        self.reset_graph()
        
    def graph_type(self) -> str:
        return GraphName.feature_graph
        
    def extract_saes(self, saes: Dict[int, List[Tuple[str, Any]]]) -> Dict[str, Any]:
        dict_saes = {}
        for layer in range(self.n_layers):
            for hook_position, sae in saes[layer]:
                dict_saes[hook_position] = sae
        return dict_saes
        
    def build_default_connection(
        self,
        seq_length: int,
        token_wise: bool = False,
        inter: bool = True,
        intra: bool = False,
    ) -> Dict:
        '''
        Sample tokens to get the shape of the activations, necessary to build the graph of features
        If token_wise, the each node is a feature of a token, else, the graph is built feature-wise
        '''
        self.reset_graph()
        
        # vision
        for layer in range(self.n_layers):
            # Setup
            self.connection[layer] = OrderedDict()
            for hook_position, _ in self.vision_saes[layer]:
                self.connection[layer][(ConnectionNode(hook_position), ConnectionIndex())] = []
                
            # Inter-layer connections
            if inter:
                for i, (hook_position_end, _) in reversed(list(enumerate(self.vision_saes[layer]))):
                    for j, (hook_position_start, _) in enumerate(self.vision_saes[layer]):
                        if i > j:
                            self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                                (ConnectionNode(hook_position_start), ConnectionIndex())
                            )
                            
            # Intra-layer connections
            if intra:
                for prev_layer in range(layer):
                    for hook_position_end, _ in self.vision_saes[layer]:
                        for hook_position_start, _ in self.vision_saes[prev_layer]:
                            self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                                (ConnectionNode(hook_position_start), ConnectionIndex())
                            )
            # Sequential connection layer-1 --> layer
            else:
                if layer == 0:
                    continue
                for hook_position_start, _ in self.vision_saes[layer-1]:
                    # connect to the first node of the next layer
                    hook_position_end = self.vision_saes[layer][0][0] 
                    self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                        (ConnectionNode(hook_position_start), ConnectionIndex())
                    )
             
        # text       
        for layer in range(self.n_layers):
            # Setup
            self.connection[layer + self.n_layers] = OrderedDict() # text has layer of: layer + n_layers
            for hook_position, _ in self.text_saes[layer]:
                self.connection[layer + self.n_layers][(ConnectionNode(hook_position), ConnectionIndex())] = []
                
            # Inter-layer connections
            if inter:
                for i, (hook_position_end, _) in reversed(list(enumerate(self.text_saes[layer]))):
                    for j, (hook_position_start, _) in enumerate(self.text_saes[layer]):
                        if i > j:
                            self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                                (ConnectionNode(hook_position_start), ConnectionIndex())
                            )
                            
            # Intra-layer connections
            if intra:
                # intra connection with vision
                for hook_position_end, _ in self.text_saes[layer]:       
                    if layer != 0 or ("crossattention.self.hook_attn_out" in hook_position_end or "mlp" in hook_position_end):
                        for vision_layer in range(self.n_layers):
                            for hook_position_start, _ in self.vision_saes[vision_layer]:
                                self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                                    (ConnectionNode(hook_position_start), ConnectionIndex())
                                )
                            
                for prev_layer in range(layer):
                    for hook_position_end, _ in self.text_saes[layer]:
                        for hook_position_start, _ in self.text_saes[prev_layer]:
                            self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                                (ConnectionNode(hook_position_start), ConnectionIndex())
                            )
                     
                        
            # Sequential connection layer-1 --> layer
            else:
                if layer == 0:
                    # connect with the last vision layer
                    hook_position_start = self.vision_saes[self.n_layers-1][-1][0]
                    hook_position_end = self.text_saes[0][0][0] 
                    if "crossattention.self.hook_attn_out" in hook_position_end or "mlp" in hook_position_end:
                        self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                            (ConnectionNode(hook_position_start), ConnectionIndex())
                    )
                else:
                    for hook_position_start, _ in self.text_saes[layer-1]:
                        # connect to the first node of the prev layer
                        hook_position_end = self.text_saes[layer][0][0] 
                        self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append(
                            (ConnectionNode(hook_position_start), ConnectionIndex())
                        )
                
        self._build_nodes(seq_length, token_wise)
        self._build_graphs()
        return self.connection
    
    def _build_nodes(self, seq_length: int, token_wise: bool) -> None:
        self.seq_length = seq_length
        self.token_wise = token_wise
        
        # text
        for sae_name, sae in self.dict_text_saes.items():
            node_shape = (self.seq_length, sae.cfg.d_sae)
            self.nodes[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct(
                t.ones(node_shape, requires_grad=False).to(self.device),
                None,
                t.ones(self.seq_length, requires_grad=False).to(self.device) if self.use_error_term else None,
            )
            self.node_scores[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct(
                t.zeros(node_shape, requires_grad=False).to(self.device),
                None,
                t.zeros(self.seq_length, requires_grad=False).to(self.device) if self.use_error_term else None,
            )
            
        # vision
        for sae_name, sae in self.dict_vision_saes.items():
            node_shape = (1, sae.cfg.d_sae) # HARDCODE
            self.nodes[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct(
                t.ones(node_shape, requires_grad=False).to(self.device),
                None,
                t.ones(1, requires_grad=False).to(self.device) if self.use_error_term else None,
            )
            self.node_scores[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct(
                t.zeros(node_shape, requires_grad=False).to(self.device),
                None,
                t.zeros(1, requires_grad=False).to(self.device) if self.use_error_term else None,
            )
            
    def _build_graphs(self) -> None:
        assert self.nodes, "Please build the nodes first"
        for _, connections in self.connection.items():
            for hook_position_end, list_hook_positions_start in connections.items():
                self.edges[hook_position_end] = OrderedDict()
                self.edge_scores[hook_position_end] = OrderedDict()
                
                for hook_position_start in list_hook_positions_start:
                    self.edges[hook_position_end][hook_position_start] = t.zeros([0]) # place holder
                    self.edge_scores[hook_position_end][hook_position_start] = t.zeros([0]) # place holder
        
    def reset_graph(self) -> None:
        self.connection: OrderedDict[int, OrderedDict[Tuple[Node, Index], List[Tuple[Node, Index]]]] = OrderedDict() # {layer: {hook_position_end: [hook_position_start]}}
        self.nodes: Dict[Tuple[Node, Index], SparseAct] = {} # {hook_position: SparseAct[act: Tensor['seq d_sae'], resc: Tensor['seq']]} 
        self.node_scores: Dict[Tuple[Node, Index], SparseAct] = {} 
        self.edges: OrderedDict[Tuple[Node, Index], OrderedDict[Tuple[Node, Index], Tensor]] = OrderedDict() # {hook_position_end: {hook_position_start: Tensor}}
        self.edge_scores: OrderedDict[Tuple[Node, Index], OrderedDict[Tuple[Node, Index], Tensor]] = OrderedDict()
        self.seq_length: int | None = None
        self.token_wise: bool | None = None
        
    def _check_graph(self) -> None:
        assert len(self.connection) > 0, "Graph is empty"
        assert len(self.nodes) > 0, "Nodes is empty"
        assert len(self.node_scores) > 0, "Node scores is empty"
        assert len(self.edges) > 0, "Edges is empty"
        assert len(self.edge_scores) > 0, "Edge scores is empty"
        assert self.seq_length is not None
        assert self.token_wise is not None
        
    def __repr__(self) -> str:
        self._check_graph()
        return nested_dict_to_string(self.connection, indent=4)
    
    def _check_feat_idx(self, feat_idx: Tuple | List) -> None:
        if self.token_wise:
            assert len(feat_idx) == 2, "Please provide the correct index"
        else:
            assert len(feat_idx) == 1, "Please provide the correct index"
            
    def _check_error_idx(self, error_idx: Tuple | List) -> None:
        if self.token_wise:
            assert len(error_idx) == 1, "Please provide the correct index"
        else:
            assert len(error_idx) == 0, "Please provide the correct index"
        
    def _active_nodes(
        self,
        sae_name: Node,
        node_idx: Index,
        reverse: bool = False,
    ) -> Tuple[List, ...]:
        self._check_graph()
        if self.token_wise:
            node = self.nodes[(sae_name, node_idx)] # act: (seq, d_sae), resc: (seq) 
        else:
            node = self.nodes[(sae_name, node_idx)].mean(dim=0) # act: (d_sae), resc: () - COLAPSED TENSOR
        
        if reverse:
            active_node = (node.act == 0).nonzero().tolist()
            active_error = (node.resc == 0).nonzero().tolist() if self.use_error_term else [] # type: ignore
        else:
            active_node = node.act.nonzero().tolist()
            active_error = node.resc.nonzero().tolist() if self.use_error_term else [] # type: ignore
            
        return active_node, active_error
            
    def active_nodes(
        self,
        sae_name: Node,
        node_idx: Index,
        reverse: bool = False,
    ) -> List[Tuple[Node, Index]]:
        active_node, active_error = self._active_nodes(sae_name, node_idx, reverse)
        
        feat_list = [
            (sae_name, FeatureIndex(idx)) 
            for idx in active_node
        ]
        error_list = [
            (sae_name, ErrorIndex(idx))
            for idx in active_error
        ]
        return feat_list + error_list # type: ignore
        
    def add_node(
        self, 
        sae_name: Node, 
        node_idx: Index,
    ) -> None:
        self._check_graph()
        pass # no implementation needed
        
    def add_edge(
        self, 
        sae_name_start: Node, 
        node_idx_start: Index, 
        sae_name_end: Node, 
        node_idx_end: Index,
    ) -> None:
        self._check_graph()
        pass # no implementation needed
    
    def delete_node(
        self, 
        sae_name: Node,
        node_idx: Index, 
    ) -> None:
        self._check_graph()
        pass # no implementation needed
    
    def delete_edge(
        self, 
        sae_name_start: Node, 
        node_idx_start: Index, 
        sae_name_end: Node, 
        node_idx_end: Index,
    ) -> None:
        self._check_graph()
        pass # no implementation needed
    
    def _check_node_shape(self, sae_name: Node, sparseact: SparseAct):
        seq = self.seq_length if "vision" not in sae_name.name else 1
        try:
            if self.token_wise:
                assert sparseact.act.shape == t.Size([seq, self.dict_saes[sae_name.name].cfg.d_sae]) # type: ignore
            else:
                assert sparseact.act.shape == t.Size([self.dict_saes[sae_name.name].cfg.d_sae])
            if self.use_error_term:
                if self.token_wise:
                    assert sparseact.resc.shape == t.Size([seq]) # type: ignore
                else:
                    assert sparseact.resc.numel() == 1 # type: ignore
        except:
            raise ValueError("Wrong input shape for nodes.")  
    
    def update_node(
        self, 
        sae_name: Node,
        node_idx: Index,  
        value_and_mask: Tuple[SparseAct, SparseAct],
    ) -> None:
        self._check_graph()
        value, mask = value_and_mask
        assert isinstance(value, SparseAct), "value is not an instance of SparseAct"
        assert isinstance(mask, SparseAct), "mask is not an instance of SparseAct"
        self._check_node_shape(sae_name, value)
        
        if not self.token_wise:
            seq = self.seq_length if "vision" not in sae_name.name else 1
            # convert to the right format of (seq, d_sae) and (seq) 
            value_act = einops.repeat(value.act, 'd_sae -> seq d_sae', seq=seq)
            mask_act = einops.repeat(mask.act, 'd_sae -> seq d_sae', seq=seq)
            value_resc = einops.repeat(value.resc, ' -> seq', seq=seq) if self.use_error_term else None
            mask_resc = einops.repeat(mask.resc, ' -> seq', seq=seq) if self.use_error_term else None
            
            value = SparseAct(value_act, None, value_resc)
            mask = SparseAct(mask_act, None, mask_resc)
            
        self.nodes[(sae_name, node_idx)] = mask.to(t.float32)
        self.node_scores[(sae_name, node_idx)] = value
    
    def update_edge(
        self, 
        sae_name_start: Node, 
        node_idx_start: Index, 
        sae_name_end: Node, 
        node_idx_end: Index,
        value_and_mask: Tuple[Tensor, Tensor],
    ) -> None:
        self._check_graph()
        value, mask = value_and_mask
        assert isinstance(value, Tensor), "value is not an instance of Tensor"
        assert isinstance(mask, Tensor), "mask is not an instance of Tensor"
        self.edges[(sae_name_end, node_idx_end)][(sae_name_start, node_idx_start)] = mask
        self.edge_scores[(sae_name_end, node_idx_end)][(sae_name_start, node_idx_start)] = value
    
    def find_deleted_nodes(self, reverse = False) -> List[Tuple[Node, Index]]:
        self._check_graph()
        
        all_list = []
        for sae_name, node_idx in self.nodes.keys():
            inactive_node, inactive_error = self._active_nodes(sae_name, node_idx, not reverse)
            feat_list = [
                (sae_name, FeatureIndex(idx))
                for idx in inactive_node
            ]
            error_list = [
                (sae_name, ErrorIndex(idx))
                for idx in inactive_error
            ]
            all_list += feat_list + error_list
        
        return all_list
    
    def find_deleted_edges(self, reverse = False) -> List[Tuple[Node, Index, Node, Index]]:
        self._check_graph()
        
        all_list = []
        for end_node, end_idx in self.edges.keys():
            for start_node, start_idx in self.edges[(end_node, end_idx)].keys():
                edge = self.edges[(end_node, end_idx)][(start_node, start_idx)]
                if not reverse:
                    inactive_node = (edge.values() == 0) # (num_active, seq, d_sae+1) or (num_active, d_sae+1)
                else:
                    inactive_node = edge.values()
                    
                d_sae_end = self.dict_saes[end_node.name].cfg.d_sae
                d_sae_start = self.dict_saes[start_node.name].cfg.d_sae
                    
                active_idx = [index for _, index in self.active_nodes(start_node, start_idx)]
                revised_active_nodes = [
                    index.idx if isinstance(index, FeatureIndex) else index.idx + (d_sae_start,) # type: ignore
                    for index in active_idx
                ]
                
                for num_active_idx in range(edge.indices().shape[1]):
                    end_node_idx = edge.indices()[:, num_active_idx].tolist()
                    # check if the end_node_idx is the error term
                    check_error = True if self.use_error_term and end_node_idx[-1] == d_sae_end else False
                    # create the end_index class to be used in the tuple
                    end_index = FeatureIndex(end_node_idx) if not check_error else ErrorIndex(end_node_idx[:-1])
                    
                    for i, idx in enumerate(revised_active_nodes):
                        if inactive_node[(num_active_idx,) + idx] > 0:
                            all_list.append(
                                (
                                    start_node, 
                                    active_idx[i], 
                                    end_node, 
                                    end_index,
                                )
                            )
        
        return all_list
    
    def iterate_nodes(self) -> List[Tuple[Node, Index]]:
        self._check_graph()
        return list(self.nodes.keys())
    
    def iterate_edges(self) -> List[Tuple[Node, Index, Node, Index]]:
        self._check_graph()
        all_edges = []
        for end, edges in self.edges.items():
            for start, edge in edges.items():
                all_edges.append(start + end)
        return all_edges
    
    def add_single_feature(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_idx: Tuple[int, ...] | List[int] | FeatureIndex,
    ) -> None:
        self._check_graph()
        self._set_feat_value(sae_name, node_idx, feat_idx, 1)
        
    def delete_single_feature(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_idx: Tuple[int, ...] | List[int] | FeatureIndex,
    ) -> None:
        self._check_graph()
        self._set_feat_value(sae_name, node_idx, feat_idx, 0)
        
    def update_single_feature(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_idx: Tuple[int, ...] | List[int] | FeatureIndex,
        value: int | float | Tensor,
    ) -> None:
        self._check_graph()
        assert isinstance(value, (int, float, Tensor)), "The value must be int, float or Tensor."
        self._set_feat_value(sae_name, node_idx, feat_idx, value)
        
    def add_single_error(
        self,
        sae_name: Node,
        node_idx: Index,
        error_idx: Tuple[int, ...] | List[int] | ErrorIndex,
    ) -> None:
        self._check_graph()
        self._set_error_value(sae_name, node_idx, error_idx, 1)
        
    def delete_single_error(
        self,
        sae_name: Node,
        node_idx: Index,
        error_idx: Tuple[int, ...] | List[int] | ErrorIndex,
    ) -> None:
        self._check_graph()
        self._set_error_value(sae_name, node_idx, error_idx, 0)
        
    def update_single_error(
        self,
        sae_name: Node,
        node_idx: Index,
        error_idx: Tuple[int, ...] | List[int] | ErrorIndex,
        value: int | float | Tensor, 
    ) -> None:
        self._check_graph()
        assert isinstance(value, (int, float, Tensor)), "The value must be int, float or Tensor."
        self._set_error_value(sae_name, node_idx, error_idx, value)
        
    def _set_feat_value(
        self,
        sae_name: Node,
        node_idx: Index,
        feat_idx: Tuple[int, ...] | List[int] | FeatureIndex,
        value: int | float | Tensor,
    ) -> None:
        self._check_graph()
        
        if isinstance(feat_idx, (tuple, list)):
            self._check_feat_idx(feat_idx)
            self.nodes[(sae_name, node_idx)].act[return_idx(feat_idx)] = value
        elif isinstance(feat_idx, FeatureIndex):
            self.nodes[(sae_name, node_idx)].act[feat_idx.as_index] = value
        else:
            raise ValueError(f"feat_idx of type {type(feat_idx)} is not supported.")
        
    def _set_error_value(
        self,
        sae_name: Node,
        node_idx: Index,
        error_idx: Tuple[int, ...] | List[int] | ErrorIndex,
        value: int | float | Tensor, 
    ) -> None:
        if not self.use_error_term:  # nothing to delete
            return
        
        if isinstance(error_idx, (tuple, list)):
            self._check_error_idx(error_idx)
            self.nodes[(sae_name, node_idx)].resc[return_idx(error_idx, 1)] = value # type: ignore
        elif isinstance(error_idx, ErrorIndex):
            self.nodes[(sae_name, node_idx)].resc[error_idx.as_index] = value # type: ignore
        else:
            raise ValueError(f"error_idx of type {type(error_idx)} is not supported.")
        
    def remove_error_term(self) -> None:
        if self.use_error_term:
            for sae_name, node_idx in self.nodes.keys():
                if self.nodes[(sae_name, node_idx)].resc is not None:
                    seq = self.seq_length if "vision" not in sae_name.name else 1
                    self.nodes[(sae_name, node_idx)].resc = t.zeros(seq, requires_grad=False).to(self.device) # type: ignore
    def forward(
        self,
        inputs: Dict[str, Tensor],
        corrupt_cache: ActivationCache | Dict[str, Tensor] | None,
        patch_deleted_comp: bool = False,
        **kwargs,
    ) -> Tuple[Tensor, Dict[str, SparseAct]]:
        '''
        Forward pass of the graph with clean tokens, if the edge exists, replace the activation with corrupted activation
        '''
        self._check_graph()
        self.model.reset_hooks()
        self.model_setup()
        
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == sae_hook_name(sae_name):
                if patch_deleted_comp and corrupt_cache is not None:
                    act_mask = self.nodes[(sae_name, (None,))].act == 0 # type: ignore
                    act[:, act_mask] = corrupt_cache[sae_hook_name(sae_name)][:, act_mask]
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif hook.name == error_term_name(sae_name) and self.use_error_term:
                if patch_deleted_comp and corrupt_cache is not None:
                    resc_mask = self.nodes[(sae_name, (None,))].resc == 0 # type: ignore
                    act[:, resc_mask] = corrupt_cache[error_term_name(sae_name)][:, resc_mask]
                fwd_cache[error_term_name(sae_name)] = act.detach()
                
            return act
        
        with t.no_grad() and self._hook_vision_sae():
            with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                with self.model.hooks(
                    fwd_hooks=[
                        (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name)) 
                        for sae_name in self.dict_saes.keys()
                    ],
                ):
                    logits = self.model(inputs)
        
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = SparseAct(
                act=fwd_cache[sae_hook_name(sae_name)],
                res=fwd_cache[error_term_name(sae_name)] if self.use_error_term else None,
            )
            
        for sae in self.dict_saes.values():
            sae.reset_hooks()
        self.model.reset_hooks()
    
        return logits, cache
    
    def __call__(self, *args, **kwargs) -> Tuple[Tensor, Dict[str, SparseAct]]:
        return self.forward(*args, **kwargs)
    
    def forward_backward_gradient(
        self,
        inputs: Dict[str, Tensor],
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        metric: Callable[[Tensor], Tensor],
        retain_graph: bool = False,
        mode: str = 'node',
        gradient_mode: str = 'standard',
        pass_through_grad: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index], SparseAct],  # node effects
        Dict[Tuple[Node, Index, Node, Index], Tensor], # edge effects
    ]:
        if mode == 'node':
            if verbose:
                print("Calculating node gradients...")
            if gradient_mode == "standard":
                node_grads, clean_cache = self._gradient_wrt_nodes(
                    inputs, metric, retain_graph, pass_through_grad, **kwargs
                )
            elif gradient_mode == "ig":
                node_grads, clean_cache = self._gradient_wrt_nodes_ig(
                    inputs, corrupt_cache, metric, retain_graph, verbose, **kwargs
                )
            else:
                raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported")
            node_effect = self._attrib_effect_node(node_grads, corrupt_cache, clean_cache)
                
            return node_effect, {}
        
        elif mode == 'edge':
            if kwargs.get('node_grads', None) is None or kwargs.get('node_effect', None) is None:
                # run the _gradient_wrt_nodes to get node_grads, and use node_effect to prune out unimportant nodes
                if verbose:
                    print("Calculating node gradients...")
                if gradient_mode == "standard":
                    node_grads, clean_cache = self._gradient_wrt_nodes(
                        inputs, metric, retain_graph, pass_through_grad, **kwargs
                    )
                elif gradient_mode == "ig":
                    node_grads, clean_cache = self._gradient_wrt_nodes_ig(
                        inputs, corrupt_cache, metric, retain_graph, verbose, **kwargs
                    )
                else:
                    raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported")
                node_effect = self._attrib_effect_node(node_grads, corrupt_cache, clean_cache)
            else:
                node_grads: Dict[Tuple[Node, Index], SparseAct] = kwargs.get('node_grads') # type: ignore
                node_effect: Dict[Tuple[Node, Index], SparseAct] = kwargs.get('node_effect') # type: ignore
                
            # Pruning
            if kwargs.get('prune', False):
                if verbose:
                    print("Pruning nodes...")
                self._prune_nodes(node_effect, verbose, **kwargs)
                
            if kwargs.get('gradient_only', False):
                if verbose:
                    print("Returning edge gradients only...")
                for name, sparse_act in node_grads.items():
                    node_grads[name] = sparse_act.to_sparse_like_self(t.ones_like(sparse_act.to_tensor()))
                    del sparse_act
                    
            edge_grads, _ = self._gradient_wrt_edges(
                inputs, corrupt_cache, node_grads, verbose, **kwargs
            )
            return node_effect, edge_grads
        
        else:
            raise NotImplementedError(f"mode {mode} is not supported")
        
    def _attrib_effect_node(
        self,
        grads: Dict,
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        clean_cache: Dict[str, SparseAct],
    ) -> Dict:
        attrib_effect = {}
        aggregate_dim = [0] if self.token_wise else [0, 1]
        
        for (node, index), grad in grads.items():
            corrupt_sparse_act = cache_to_sparseact(
                corrupt_cache, 
                sae_hook_name(node.name), 
                error_term_name(node.name) if self.use_error_term else None,
            )
            attrib_effect[(node, index)] = ( # act: (seq, d_sae), resc: (seq) || act: (d_sae), resc: ()
                grad @
                (corrupt_sparse_act - clean_cache[node.name])
            ).sum(aggregate_dim)
            
        return attrib_effect
    
    def _gradient_wrt_nodes(
        self,
        inputs: Dict[str, Tensor],
        metric: Callable[[Tensor], Tensor],
        retain_graph: bool = False,
        pass_through_grad: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index], SparseAct],  # node effects
        Dict[str, SparseAct]
    ]:
        '''
        Forward pass of the graph with clean tokens, if the edge exists, replace the activation with corrupted activation
        Backward pass on the graph wrt the metric
        Return the gradients wrt nodes, edges, and activation cache
        '''
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()

        bwd_cache = {}
        pass_through_cache = {}
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == sae_hook_name(sae_name):
                bwd_cache[sae_hook_name(sae_name)] = grad.detach()
                
            elif hook.name == output_hook_name(sae_name):
                if self.use_error_term:
                    # we have: output = recon + stop_grad(error_term)
                    # so, the TRUE error_grad (if not stop_grad) is output_grad 
                    # due to: output = recon + error_term
                    if "vision" not in sae_name:
                        bwd_cache[error_term_name(sae_name)] = grad.detach()
                    else:
                        bwd_cache[error_term_name(sae_name)] = t.zeros_like(grad).mean(dim=1, keepdim=True)
                if pass_through_grad:
                    pass_through_cache[output_hook_name(sae_name)] = grad.detach()
                    
            elif hook.name == input_hook_name(sae_name):
                if "vision" in sae_name:
                    grad.zero_()
                else:
                    if pass_through_grad:
                        # we have to modify inplace instead of grad = ... and then return grad
                        # because, returning a tensor in bwd pass hook is buggy somehow
                        grad.copy_(pass_through_cache[output_hook_name(sae_name)])
                    
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == sae_hook_name(sae_name):
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif error_term_name(sae_name) == hook.name and self.use_error_term:
                fwd_cache[error_term_name(sae_name)] = act.detach()
                    
            return act
        
        with t.set_grad_enabled(True):
            with self._detach_error_term(True) and self._hook_vision_sae():
                with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                    with self.model.hooks(
                        fwd_hooks=[
                            (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name)) 
                            for sae_name in self.dict_saes.keys()
                        ],
                        bwd_hooks=[
                            (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_bwd, sae_name=sae_name)) 
                            for sae_name in self.dict_saes.keys()
                        ],
                    ):
                        metric(self.model(inputs)).backward(retain_graph=retain_graph)
                
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = cache_to_sparseact(
                bwd_cache,
                sae_hook_name(node.name),
                error_term_name(node.name) if self.use_error_term else None,
            )
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                fwd_cache,
                sae_hook_name(sae_name),
                error_term_name(sae_name) if self.use_error_term else None,
            )
            
        self.model.reset_hooks()    
        for sae in self.dict_saes.values():
            sae.reset_hooks()       
        return node_grads, cache
    
    def _gradient_wrt_nodes_ig(
        self,
        inputs: Dict[str, Tensor],
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        metric: Callable[[Tensor], Tensor],
        retain_graph: bool = False,
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index], SparseAct],  # node effects
        Dict[str, SparseAct]
    ]:
        steps = kwargs.get('steps', 10)
        
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()

        bwd_cache: Dict[str, Tensor] = {}
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == sae_hook_name(sae_name):
                create_list(bwd_cache, sae_hook_name(sae_name))
                bwd_cache[sae_hook_name(sae_name)] += grad.detach()
                
            elif hook.name == output_hook_name(sae_name):
                if self.use_error_term:
                    # we have: output = recon + stop_grad(error_term)
                    # so, the TRUE error_grad (if not stop_grad) is output_grad 
                    # due to: output = recon + error_term
                    create_list(bwd_cache, error_term_name(sae_name))
                    if "vision" not in sae_name:
                        bwd_cache[error_term_name(sae_name)] += grad.detach()
                    else:
                        bwd_cache[error_term_name(sae_name)] += t.zeros_like(grad).mean(dim=1, keepdim=True)
                    
        fwd_cache = {}
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, target_name: str, frac: float) -> Tensor:
            if hook.name == sae_hook_name(sae_name):
                # interpolate for integrated gradients
                if hook.name == sae_hook_name(target_name):
                    act = interpolate(
                        corrupt_cache[sae_hook_name(sae_name)], 
                        act, 
                        frac,
                    )
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif error_term_name(sae_name) == hook.name and self.use_error_term:
                # interpolate for integrated gradients
                if hook.name == error_term_name(target_name):
                    act = interpolate(
                        corrupt_cache[error_term_name(sae_name)], 
                        act, 
                        frac,
                    )
                fwd_cache[error_term_name(sae_name)] = act.detach()
                    
            return act
        
        with t.set_grad_enabled(True):
            with self._detach_error_term(True) and self._hook_vision_sae():
                with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                    for target_name in self.dict_saes.keys():
                        for step in range(steps):
                            frac = step / steps
                            with self.model.hooks(
                                fwd_hooks=[
                                    (
                                        lambda name, sae_name=sae_name: sae_name in name, 
                                        partial(hook_sae_fwd, sae_name=sae_name, target_name=target_name, frac=frac)
                                    ) for sae_name in self.dict_saes.keys()
                                ],
                                bwd_hooks=[
                                    (
                                        lambda name, target_name=target_name: target_name in name, 
                                        partial(hook_sae_bwd, sae_name=target_name)
                                    ) 
                                ],
                            ):
                                metric(self.model(inputs)).backward(retain_graph=retain_graph)
               
        # average the gradients                 
        for key in bwd_cache.keys():
            bwd_cache[key] /= steps
                
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = cache_to_sparseact(
                bwd_cache,
                sae_hook_name(node.name),
                error_term_name(node.name) if self.use_error_term else None,
            )
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                fwd_cache,
                sae_hook_name(sae_name),
                error_term_name(sae_name) if self.use_error_term else None,
            )
            
        self.model.reset_hooks()    
        for sae in self.dict_saes.values():
            sae.reset_hooks()       
        return node_grads, cache
        
    def _gradient_wrt_edges(
        self,
        inputs: Dict[str, Tensor],
        corrupt_cache: ActivationCache | Dict[str, Tensor],
        node_grads: Dict[Tuple[Node, Index], SparseAct],
        verbose: bool = False,
        **kwargs,
    ) -> Tuple[
        Dict[Tuple[Node, Index, Node, Index], Tensor], # edge effects
        Dict[str, SparseAct]
    ]:
        
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()
            
        gradient_mode = kwargs.get('edge_gradient_mode', 'gradient')
            
        _, clean_cache = self.forward(inputs, corrupt_cache=None)
                    
        edge_grads: Dict[Tuple[Node, Index, Node, Index], Tensor] = {}
        for layer, connection in tqdm(self.connection.items(), disable=not verbose):
            if verbose:
                print(f"Layer {layer}:")
            for hook_position_end, list_hook_positions_start in tqdm(connection.items(), disable=not verbose):
                assert hook_position_end in node_grads, f"Node gradient of {hook_position_end} is not provided."
                
                for hook_position_start in list_hook_positions_start:
                    corrupt_sparse_act = cache_to_sparseact(
                        corrupt_cache, 
                        sae_hook_name(hook_position_start[0].name), 
                        error_term_name(hook_position_start[0].name) if self.use_error_term else None,
                    )
                    right_vec = corrupt_sparse_act - clean_cache[hook_position_start[0].name]
                    
                    if gradient_mode == "gradient":
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution(
                            inputs,
                            hook_position_end,
                            hook_position_start,
                            node_grads[hook_position_end],
                            right_vec,
                            layer,
                            **kwargs,
                        )
                    else:
                        raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported")
                                  
        return edge_grads, clean_cache
    
    def _edge_attribution(
        self,
        inputs: Dict[str, Tensor],
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        layer: int,
        **kwargs,
    ) -> Tensor:
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str, bwd_cache: Dict) -> None:
            if hook.name == sae_hook_name(hook_position_start[0].name):
                # only store the gradient at the start position
                bwd_cache[sae_hook_name(hook_position_start[0].name)] = grad.detach()
                
            elif hook.name == output_hook_name(hook_position_start[0].name):
                # we have: output = recon + stop_grad(error_term)
                # so, the TRUE error_grad (if not stop_grad) is output_grad 
                # due to: output = recon + error_term
                if self.use_error_term:
                    # only store the gradient at the start position
                    if "vision" not in sae_name:
                        bwd_cache[error_term_name(hook_position_start[0].name)] = grad.detach()
                    else:
                        bwd_cache[error_term_name(hook_position_start[0].name)] = t.zeros_like(grad).mean(dim=1, keepdim=True)

            # IMPORTANT NOTE for reproducibility:
            # we zero grad of intermediate components
            # but zero grad the output hook of SAE will ALSO ZERO GRAD the resid mid / post --> no downstream grads
            # so, we instead zero grad of the input hook of SAE 
            elif hook.name == input_hook_name(sae_name):
                if hook.name != input_hook_name(hook_position_end[0].name):
                    grad.zero_()

        def hook_sae_fwd(act: Tensor, hook: HookPoint, to_bwd_cache: Dict) -> Tensor:
            if hook.name == sae_hook_name(hook_position_end[0].name):
                # store activation for backward later
                to_bwd_cache[sae_hook_name(hook_position_end[0].name)] = act
                
            elif error_term_name(hook_position_end[0].name) == hook.name and self.use_error_term:
                # store activation for backward later
                to_bwd_cache[error_term_name(hook_position_end[0].name)] = act
                    
            return act
        
        d_sae_end = self.dict_saes[hook_position_end[0].name].cfg.d_sae
        d_sae_start = self.dict_saes[hook_position_start[0].name].cfg.d_sae
        
        to_bwd_cache = {}
        bwd_cache = {}
        edge_effect = OrderedDict()
        with t.set_grad_enabled(True):
            with self._detach_error_term(False, hook_position_end[0].name) and self._hook_vision_sae():
                with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term):
                    with self.model.hooks(
                        fwd_hooks=[
                            (lambda name: True, partial(
                                hook_sae_fwd, 
                                to_bwd_cache=to_bwd_cache, 
                            ))
                        ],
                        bwd_hooks=[
                            (lambda name, sae_name=sae_name: sae_name in name, partial(
                                hook_sae_bwd, 
                                sae_name=sae_name, 
                                bwd_cache=bwd_cache,
                            )) for sae_name in self.dict_saes.keys()
                        ],
                    ):
                        self.model.forward(inputs)
                        
                        aggregate_dim = [0] if self.token_wise else [0, 1]
                        to_bwd = ( # (seq, d_sae+1) || (d_sae+1)
                            cache_to_sparseact(
                                to_bwd_cache, 
                                sae_hook_name(hook_position_end[0].name), 
                                error_term_name(hook_position_end[0].name) if self.use_error_term else None,
                            ) @ leftvec.detach()
                        ).sum(aggregate_dim).to_tensor()
                        del to_bwd_cache
                        
                        for active_idx, (end_node, end_index) in enumerate(self.active_nodes(*hook_position_end)):
                            if isinstance(end_index, ErrorIndex):
                                # the last index is error: shape (d_sae+1) so last index is d_sae
                                to_bwd[end_index.idx + (d_sae_end,)].backward(retain_graph=True)
                                index = t.tensor(list(end_index.idx + (d_sae_end,)), device=self.device)
                            elif isinstance(end_index, FeatureIndex):
                                to_bwd[end_index.idx].backward(retain_graph=True)
                                index = t.tensor(list(end_index.idx), device=self.device)
                            else:
                                raise ValueError(f"end_index of type {type(end_index)} is not supported.")
                            '''
                            edge_effect shape (seq, d_sae+1, seq, d_sae+1) or (d_sae+1, d_sae+1) in sparse_coo tensor
                            
                            the sparse_coo will have the shape:
                            --> indices of shape (2, num_active) or (1, num_active)
                            --> values of shape (num_active, seq, d_sae+1) or (num_active, d_sae+1)
                            '''
                            edge_effect[active_idx] = (
                                index,
                                ( # (seq, d_sae+1) || (d_sae+1)
                                    cache_to_sparseact(
                                        bwd_cache, 
                                        sae_hook_name(hook_position_start[0].name), 
                                        error_term_name(hook_position_start[0].name) if self.use_error_term else None,
                                    ) @ rightvec
                                ).sum(aggregate_dim).to_tensor()
                            )
                        
                        del bwd_cache
                    
        seq_start = rightvec.act.shape[1] # type: ignore
        seq_end = leftvec.act.shape[1] # type: ignore
        num_end = d_sae_end
        num_start = d_sae_start
        if self.use_error_term:
            num_end += 1
            num_start += 1
        
        if len(edge_effect.keys()) != 0:
            indices = t.stack([val[0] for val in edge_effect.values()], dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([val[1] for val in edge_effect.values()], dim=0) # shape (num_active, seq, d_sae+1) or (num_active, d_sae+1)
        # if no active nodes, return empty tensor
        else:
            indices = t.empty((2, 0) if self.token_wise else (1, 0), dtype=t.long).to(self.device)
            values = t.empty((0, seq_start, num_start) if self.token_wise else (0, num_start), dtype=t.float).to(self.device)

        if self.token_wise:
            return t.sparse_coo_tensor(indices, values, size=(seq_end, num_end, seq_start, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()
        
    def _prune_nodes(
        self,
        node_effect: Dict[Tuple[Node, Index], SparseAct],
        verbose: bool = False,
        **kwargs,
    ) -> None:
        if kwargs.get("node_threshold", None) is None:
            raise ValueError("Please provide the node_threshold")
        else:
            node_threshold = kwargs.get("node_threshold")
            assert isinstance(node_threshold, (float, int, Tensor)), "node_threshold must be a int, float, or Tensor"
            
            for node, index in tqdm(node_effect.keys(), disable=not verbose):
                if kwargs.get("reverse_prune_node", False):
                    effect_feat_mask = node_effect[(node, index)].act.abs() < node_threshold
                else:
                    effect_feat_mask = node_effect[(node, index)].act.abs() > node_threshold
                    
                if verbose:
                    num_pruned = (~effect_feat_mask).sum().item()
                    print(f"{(node, index)}: pruned ({num_pruned / effect_feat_mask.numel() * 100:.5f}%) features")
                
                if self.use_error_term:
                    if kwargs.get("reverse_prune_node", False):
                        effect_error_mask = node_effect[(node, index)].resc.abs() < node_threshold # type: ignore
                    else:
                        effect_error_mask = node_effect[(node, index)].resc.abs() > node_threshold # type: ignore
                        
                    if verbose:
                        num_pruned = (~effect_error_mask).sum().item()
                        print(f"{(node, index)}: pruned ({num_pruned / effect_error_mask.numel() * 100:.5f}%) errors")
                        
                # replace the nodes with mask to prune
                mask = SparseAct(
                    act=effect_feat_mask,
                    resc=effect_error_mask if self.use_error_term else None, # type: ignore
                )
                self.nodes[(node, index)] = self.nodes[(node, index)] * mask
            
    def model_setup(self):
        pass
    
    def run_model(
        self, 
        inputs: Dict[str, Tensor], 
        use_error_term: bool | None = None,
    ) -> Tuple[Tensor, ActivationCache]:
        
        with self._hook_vision_sae():
            out = self.model.run_with_cache_with_saes(
                inputs,
                saes=self._saes_to_list(),
                use_error_term=self.use_error_term if use_error_term is None else use_error_term,
                names_filter=lambda name: "sae" in name,
            )
        return out # type: ignore
    
    def _saes_to_list(self) -> List[Any]:
        return [sae for _, sae in self.dict_saes.items()]

    @contextmanager
    def _detach_error_term(self, detach: bool, sae_name: str | None = None):
        orig_detach_error_term = {}
        orig_disable_error_grad = {}
        try:
            for name, sae in self.dict_saes.items():
                if sae_name is None or sae_name == name:
                    orig_detach_error_term[name] = sae.detach_error_term
                    sae.detach_error_term = detach
                else:
                    orig_detach_error_term[name] = sae.detach_error_term
                    sae.detach_error_term = True # default detach error term
                
                orig_disable_error_grad[name] = sae.disable_error_grad
                sae.disable_error_grad = False # allow grad flows through error hook
            yield
        finally:
            for name, sae in self.dict_saes.items():
                sae.detach_error_term = orig_detach_error_term[name]
                sae.disable_error_grad = orig_disable_error_grad[name]
                
    @contextmanager
    def _hook_vision_sae(self):
        pass_through_vision_sae_cache = [t.tensor(0)] # placeholder
        def hook_fn(act: Tensor, hook: HookPoint, sae_name: str):
            if input_hook_name(sae_name) == hook.name:
                pass_through_vision_sae_cache[0] = act # nodetach allow gradient to flow
                return act.mean(dim=1, keepdim=True)
                
            elif output_hook_name(sae_name) == hook.name:
                return pass_through_vision_sae_cache[0] + (act - act.detach()) / 2 # allow gradient to flow through sae
            
            elif error_term_name(sae_name) == hook.name:
                act = t.zeros_like(act).mean(dim=1, keepdim=True)
                return act
        
        try: 
            for sae_name, vision_sae in self.dict_vision_saes.items():
                vision_sae.add_hook(
                    lambda name: True,
                    partial(hook_fn, sae_name=sae_name),
                    dir="fwd",
                )
            yield
        finally:
            for vision_sae in self.dict_vision_saes.values():
                vision_sae.reset_hooks()