extra_materials/graph/Feature_Graph_Trans.py

from tracemalloc import start
from typing import Literal
from .Feature_Graph import *
from transformer_lens.utils import get_act_name

def add_cache(all_cache: Dict[str, float | Tensor], cache: Dict[str, Tensor]):
    for key, value in cache.items():
        if key not in all_cache:
            all_cache[key] = 0
        all_cache[key] += value
        
def _hook_name(sae_name: str, name: str | None):
    return sae_name + "." + str(name)

@t.no_grad()
def get_ln_constant(
    cache: Dict[str, Tensor] | ActivationCache, 
    vector: Tensor, # (1, 1, d_model) | (b, 1, d_model) | (b, seq, d_model)
    layer: int, 
    pos: int | None, 
    ln: Literal['ln1', 'ln2'] = 'ln2',
    eps: float = 1e-6,
) -> Tensor:
    assert ln in ['ln1', 'ln2'], f"ln must be either 'ln1' or 'ln2', got {ln}"
    x_act_name = get_act_name('resid_mid', layer) if ln == "ln2" else get_act_name('resid_pre', layer)
    y_act_name = get_act_name('normalized', layer, ln)

    if pos is not None:
        x = cache[x_act_name][:, pos].unsqueeze(1) # (b, 1, d_model)
        y = cache[y_act_name][:, pos].unsqueeze(1) # (b, 1, d_model)
        
        if vector.squeeze(1).ndim == 3:
            vector = vector[:, pos].unsqueeze(1) # (b, 1, d_model)
            
        denom = t.sum(x * vector, dim=-1, keepdim=True) # (b, 1, 1)
        numer = t.sum(y * vector, dim=-1, keepdim=True) # (b, 1, 1)
    else:
        x = cache[x_act_name] # (b, seq, d_model)
        y = cache[y_act_name] # (b, seq, d_model)
        
        denom = t.sum(x * vector, dim=-1, keepdim=True) # (b, seq, 1)
        numer = t.sum(y * vector, dim=-1, keepdim=True) # (b, seq, 1)

    # Handle zero division
    mask = t.abs(denom) < eps
    safe_denom = t.where(mask, t.ones_like(denom), denom)
    result = numer / safe_denom
    
    # Set result to 0 where denominator is zero
    return t.where(mask, t.zeros_like(result), result) * vector # (b, 1 or seq, d_model)

@t.no_grad()
def get_attn_head_contribs(
    vector: Tensor, # (1, 1, d_model) | (b, 1, d_model) | (b, seq, d_model)
    model: HookedSAETransformer, 
    cache: Dict[str, Tensor] | ActivationCache, 
    layer: int,
    dst_pos: int | None,
    sum_head: bool = True,
):
    attn_pattern = cache[get_act_name('pattern', layer)] # (b, h, dst, src)
    batch_size = attn_pattern.shape[0]
    
    if vector.shape[0] != batch_size: # upsample to batchsize
        vector = vector.repeat((batch_size, 1, 1)) # (b, 1 or seq, d_model)
    
    if dst_pos is not None:
        attn_pattern = attn_pattern[:, :, dst_pos] # (b, h, src)
        
        if vector.squeeze(1).ndim == 3:
            vector = vector[:, dst_pos].unsqueeze(1) # (b, 1, d_model2)
        
        grad_outs = einops.einsum(
            model.W_O[layer], vector.squeeze(1),
            'h d_head d_model2, b d_model2 -> b h d_head',
        )
    
        grad_vals = einops.einsum(
            attn_pattern, grad_outs,
            'b h src, b h d_head -> b h src d_head',
        )
    
    else:
        grad_outs = einops.einsum(
            model.W_O[layer], vector,
            'h d_head d_model2, b dst d_model2 -> b h dst d_head',
        )
        
        grad_vals = einops.einsum(
            attn_pattern, grad_outs,
            'b h dst src, b h dst d_head -> b h src d_head',
        )
        
    grad = einops.einsum(
        model.W_V[layer], grad_vals,
        'h d_model1 d_head, b h src d_head -> b h src d_model1'
    )
 
    if sum_head:
        grad = t.sum(grad, dim=1) # (b, seq, d_model)

    return grad # (b, seq, d_model) | (b, h. seq, d_model)

def _check_shape(vec: Tensor, pos: int | None):
    if vec.squeeze(1).ndim != 3:
        assert pos is not None
    else:
        assert pos is None

@t.no_grad()
def gradient_ln_only(
    end_feature_vec: Tensor, # (1, 1, d_model) | (b, 1, d_model) | (b, seq, d_model)
    start_feature_vec: Tensor | None, # (d_sae, d_model)
    layer_end: int,
    pos_end: int | None,
    seq_length: int,
    batch_size: int,
    use_error_term: bool,
    cache: Dict[str, Tensor] | ActivationCache,
    device: str | t.device,
    ln: Literal["ln1", "ln2"] | None = None,
) -> SparseAct:
    d_model = end_feature_vec.shape[-1]
    
    _check_shape(end_feature_vec, pos_end)
    
    if ln is not None:
        feature_scale = get_ln_constant( # (b, 1 or seq, d_model)
            cache,
            end_feature_vec,
            layer_end,
            pos_end,
            ln=ln,
        )
    else:
        feature_scale = end_feature_vec
        
    if pos_end is not None:
        grad = t.zeros((batch_size, seq_length, d_model), device=device) # (b, seq, d_model)
        grad[:, pos_end] = feature_scale.squeeze(1) # (b, d_model)
    else:
        grad = feature_scale # (b, seq, d_model)
        
    if start_feature_vec is not None:
        return SparseAct(
            act=grad @ start_feature_vec.T, # (b, seq, d_sae)
            res=grad if use_error_term else None, # (b, seq, d_model)
            resc=None,
        )
    else:
        return SparseAct(
            act=grad, # (b, seq, d_model)
            res=grad if use_error_term else None, # (b, seq, d_model)
            resc=None,
        )
        
@t.no_grad()
def gradient_with_mlp(
    end_feature_vec: Tensor, # (1, 1, d_model) | (b, 1, d_model) | (b, seq, d_model)
    start_feature_vec: Tensor | None, # (d_sae, d_model)
    transcoder_enc: Tensor, # (d_model, d_sae)
    transcoder_dec: Tensor, # (d_sae, d_model)
    transcoder_act: Tensor, # (b, seq, d_sae)
    layer_end: int,
    pos_end: int | None,
    seq_length: int,
    batch_size: int,
    use_error_term: bool,
    cache: Dict[str, Tensor] | ActivationCache,
    device: str | t.device,
) -> SparseAct:
    d_model = end_feature_vec.shape[-1]
    _check_shape(end_feature_vec, pos_end)
    
    if pos_end is not None:
        # [(1 or b, 1, d_model) @ (d_model, d_sae)] * (b, 1, d_sae) -> (b, 1, d_sae)
        grad_post = (end_feature_vec @ transcoder_dec.T) * transcoder_act[:, pos_end].unsqueeze(1)
        grad_enc = grad_post @ transcoder_enc.T # (b, 1, d_model)
        
        feature_scale = get_ln_constant( # (b, 1, d_model)
            cache,
            grad_enc,
            layer_end,
            pos_end,
            ln="ln2",
        )
        grad = t.zeros((batch_size, seq_length, d_model), device=device) # (b, seq, d_model)
        grad[:, pos_end] = feature_scale.squeeze(1) # (b, d_model)
    else:
        # [(b, seq, d_model) @ (d_model, d_sae)] * (b, seq, d_sae) -> (b, seq, d_sae)
        grad_post = (end_feature_vec @ transcoder_dec.T) * transcoder_act
        grad_enc = grad_post @ transcoder_enc.T # (b, seq, d_model)
        
        grad = get_ln_constant( # (b, seq, d_model)
            cache,
            grad_enc,
            layer_end,
            pos_end,
            ln="ln2",
        )
    
    if start_feature_vec is not None:
        return SparseAct(
            act=grad @ start_feature_vec.T, # (b, seq, d_sae)
            res=grad if use_error_term else None, # (b, seq, d_model)
            resc=None,
        )
    else:
        return SparseAct(
            act=grad, # (b, seq, d_model)
            res=grad if use_error_term else None, # (b, seq, d_model)
            resc=None,
        )

@t.no_grad()
def gradient_with_attn(
    model: HookedSAETransformer,
    end_feature_vec: Tensor, # (1, 1, d_model) | (b, 1, d_model) | (b, seq, d_model)
    start_feature_vec: Tensor | None, # (d_sae, d_model)
    layer_end: int,
    pos_end: int | None,
    use_error_term: bool,
    cache: Dict[str, Tensor] | ActivationCache,
    device: str | t.device,
) -> SparseAct:
    _check_shape(end_feature_vec, pos_end)
    grad_through_attn = get_attn_head_contribs(end_feature_vec, model, cache, layer_end, pos_end) # (b, seq, d_model)
    
    grad = get_ln_constant( # (b, seq, d_model)
        cache,
        grad_through_attn,
        layer_end,
        None,
        ln="ln1",
    )
    
    if start_feature_vec is not None:
        return SparseAct(
            act=grad @ start_feature_vec.T, # (b, seq, d_sae)
            res=grad if use_error_term else None, # (b, seq, d_model)
            resc=None,
        )
    else:
        return SparseAct(
            act=grad, # (b, seq, d_model)
            res=grad if use_error_term else None, # (b, seq, d_model)
            resc=None,
        )


class Feature_Graph_Trans(Feature_Graph):
    def __init__(
        self,
        model: HookedSAETransformer,
        saes: Dict[int, List[Tuple[str, Any]]], # {layer: list[{hook_position: HookedSAE}]}, can define granularity here
        use_error_term: bool = False,
    ):
        super().__init__(model, saes, use_error_term)
        
        self.process_transcoder()
                
    def process_transcoder(self):
        self.input_hooks = []
        self.output_hooks = []
        self.transcoders = []
        self.non_transcoders = []
        check_mlp_out = False
        self.check_attn_out = False
        self.check_resid_pre = False
        for sae in self.dict_saes.values():
            sae = sae.to(self.device)
            if getattr(sae, 'input_hook', False) and getattr(sae, 'output_hook', False):
                self.input_hooks.append(sae.input_hook)
                self.output_hooks.append(sae.output_hook)
                self.transcoders.append(sae)
                
                if "mlp_out" in sae.output_hook:
                    check_mlp_out = True
            else:
                self.non_transcoders.append(sae)
                
            if "attn_out" in sae.cfg.hook_name:
                self.check_attn_out = True
            if "resid_pre" in sae.cfg.hook_name:
                self.check_resid_pre = True
            
        assert check_mlp_out, "Transcoder needs to be provided at the mlp_out hook."
                
    def forward(
        self,
        clean_token: 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 = {}
        
        with t.no_grad():
            with self._setup_forward_model_hook(transfer_grad=False):
                with self._setup_fwd_sae_hook(
                    fwd_cache=fwd_cache, corrupt_cache=corrupt_cache, patch_deleted_comp=patch_deleted_comp
                ):
                    logits = self.model(clean_token)
        
        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,
            )
            
        for sae in self.dict_saes.values():
            sae.reset_hooks()
        self.model.reset_hooks()
        
        return logits, cache
    
    def forward_backward_gradient(
        self,
        clean_token: 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(
                    clean_token, metric, retain_graph, pass_through_grad, **kwargs
                )
            elif gradient_mode == "ig":
                node_grads, clean_cache = self._gradient_wrt_nodes_ig(
                    clean_token, corrupt_cache, metric, retain_graph, verbose, **kwargs
                )
            elif gradient_mode == "virtual_weight":
                node_grads, clean_cache = self._gradient_wrt_nodes_vw(
                    clean_token, metric, retain_graph, pass_through_grad, **kwargs
                )
            elif gradient_mode == "virtual_weight_ig":
                node_grads, clean_cache = self._gradient_wrt_nodes_vw_ig(
                    clean_token, corrupt_cache, metric, retain_graph, **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(
                        clean_token, metric, retain_graph, pass_through_grad, **kwargs
                    )
                elif gradient_mode == "ig":
                    node_grads, clean_cache = self._gradient_wrt_nodes_ig(
                        clean_token, corrupt_cache, metric, retain_graph, verbose, **kwargs
                    )
                elif gradient_mode == "virtual_weight":
                    node_grads, clean_cache = self._gradient_wrt_nodes_vw(
                        clean_token, metric, retain_graph, pass_through_grad, **kwargs
                    )
                elif gradient_mode == "virtual_weight_ig":
                    node_grads, clean_cache = self._gradient_wrt_nodes_vw_ig(
                        clean_token, corrupt_cache, metric, retain_graph, **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(
                clean_token, corrupt_cache, node_grads, verbose, **kwargs
            )
            return node_effect, edge_grads
        else:
            raise NotImplementedError(f"mode {mode} is not supported")
    
    def _gradient_wrt_nodes(
        self,
        clean_token: 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]
    ]:
        self._check_graph()
        self.model_setup()
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()
            
        fwd_cache = {}
        bwd_cache = {}
            
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self._setup_forward_model_hook(transfer_grad=kwargs.get("transfer_grad", True)):
                    with self._setup_fwd_bwd_grad_sae_hook(
                        fwd_cache=fwd_cache, 
                        bwd_cache=bwd_cache, 
                        pass_through_grad=pass_through_grad,
                    ):
                        metric(self.model(clean_token)).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,
        clean_token: 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()

        fwd_cache = {}
        bwd_cache = {}
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self._setup_forward_model_hook(transfer_grad=kwargs.get("transfer_grad", True)):
                    for target_name in self.dict_saes.keys():
                        for step in range(steps):
                            frac = step / steps
                            with self._setup_fwd_bwd_grad_sae_hook_ig(
                                target_name=target_name,
                                frac=frac,
                                fwd_cache=fwd_cache,
                                bwd_cache=bwd_cache,
                                corrupt_cache=corrupt_cache,
                            ):
                                metric(self.model(clean_token)).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_nodes_vw(
        self,
        clean_token: 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]
    ]:
        '''
        Using virtual weight to compute node grad
        '''
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
            
        # unpatched cache on clean token
        _, unpatch_clean_cache = self.model.run_with_cache(clean_token)
        
        sink_hook_name = get_act_name("resid_post", self.n_layers-1) # last layer resid post hook
        
        sink_node_cache: Dict[str, Tensor] = {}
        def hook_bwd(tens: Tensor, hook: HookPoint):
            sink_node_cache[hook.name] = tens.detach() # type: ignore
            
            
        fwd_cache = {}
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self.model.hooks(
                    bwd_hooks=[(sink_hook_name, hook_bwd)]
                ):
                    with self._setup_forward_model_hook(transfer_grad=False): # save memory, we don't need transfer_grad here
                        with self._setup_fwd_bwd_grad_sae_hook(
                            fwd_cache=fwd_cache, 
                            bwd_cache={}, 
                            pass_through_grad=pass_through_grad,
                        ):
                            metric(self.model(clean_token)).backward(retain_graph=retain_graph)
                        
        current_grad = sink_node_cache[sink_hook_name] # (b, seq, d_model)
        
        bwd_cache = self._TE_using_virtual_weight(
            current_grad,
            fwd_cache=fwd_cache,
            unpatch_clean_cache=unpatch_clean_cache,
        )
            
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = bwd_cache[node.name]
            
        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_vw_ig(
        self,
        clean_token: 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]
    ]:
        '''
        Using virtual weight to compute node grad
        '''
        steps = kwargs.get("steps", 10)
        
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        
        sink_hook_name = get_act_name("resid_post", self.n_layers-1) # last layer resid post hook
        
        sink_node_cache: Dict[str, Tensor] = {}
        def hook_bwd(tens: Tensor, hook: HookPoint):
            sink_node_cache[hook.name] = tens.detach() # type: ignore
        
        all_fwd_cache = {}
        all_bwd_cache = {}
        for target_name in self.dict_saes.keys():
            for step in range(steps):
                frac = step / steps

                fwd_cache = {}
                with t.set_grad_enabled(True):
                    with self._detach_error_term(True):
                        with self.model.hooks(
                            bwd_hooks=[(sink_hook_name, hook_bwd)]
                        ):
                            with self._setup_forward_model_hook(transfer_grad=False): # save memory, we don't need transfer_grad here
                                with self._setup_fwd_bwd_grad_sae_hook_ig(
                                    target_name=target_name,
                                    frac=frac,
                                    fwd_cache=fwd_cache, 
                                    bwd_cache={}, 
                                    corrupt_cache=corrupt_cache, 
                                ):
                                    metric(self.model(clean_token)).backward(retain_graph=retain_graph)
                                    
                # run the model on interpolate intervention
                with self._setup_forward_model_hook(transfer_grad=False): # save memory, we don't need transfer_grad here
                    with self._setup_virtual_weight_sae_hook_ig(
                        target_name=target_name,
                        frac=frac,
                        corrupt_cache=corrupt_cache, 
                    ):
                        _, unpatch_clean_cache = self.model.run_with_cache(clean_token)   
                                
                current_grad = sink_node_cache[sink_hook_name] # (b, seq, d_model)
                
                bwd_cache = self._TE_using_virtual_weight(
                    current_grad,
                    fwd_cache=fwd_cache,
                    unpatch_clean_cache=unpatch_clean_cache,
                )
                
                if step == 0:
                    add_cache(all_fwd_cache, fwd_cache)
                
                add_cache(all_bwd_cache, bwd_cache)
                
        # average the gradients                 
        for key in all_bwd_cache.keys():
            all_bwd_cache[key] /= steps
            
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = all_bwd_cache[node.name]
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                all_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 _TE_using_virtual_weight(
        self,
        current_grad: Tensor,
        fwd_cache: ActivationCache | Dict[str, Tensor],
        unpatch_clean_cache: ActivationCache | Dict[str, Tensor],
    ) -> Dict[str, Tensor]:
        bwd_cache = {}
        for layer in reversed(range(self.n_layers)):
            mlp_name = get_act_name("mlp_out", layer)
            attn_name = get_act_name("attn_out", layer)
            resid_pre_name = get_act_name("resid_pre", layer)
            
            # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
            bwd_cache[mlp_name] = SparseAct(
                act=current_grad @ self.dict_saes[mlp_name].W_dec.T,
                res = current_grad if self.use_error_term else None,
            )
            
            grad_through_mlp = gradient_with_mlp(
                end_feature_vec=current_grad, # (b, seq, d_model)
                start_feature_vec=None,
                transcoder_enc=self.dict_saes[mlp_name].W_enc, # (d_model, d_sae)
                transcoder_dec=self.dict_saes[mlp_name].W_dec, # (d_sae, d_model)
                transcoder_act=fwd_cache[sae_hook_name(mlp_name)], # (b, seq, d_sae)
                layer_end=layer,
                pos_end=None, 
                seq_length=self.seq_length, # type: ignore
                batch_size=current_grad.shape[0],
                use_error_term=self.use_error_term,
                cache=unpatch_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
            
            current_grad = current_grad + grad_through_mlp # resid_mid_grad
            
            if self.check_attn_out:
                # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                bwd_cache[attn_name] = SparseAct(
                    act=current_grad @ self.dict_saes[attn_name].W_dec.T,
                    res = current_grad if self.use_error_term else None,
                )
            
            grad_through_attn = gradient_with_attn(
                model=self.model,
                end_feature_vec=current_grad, # (b, seq, d_model)
                start_feature_vec=None,
                layer_end=layer,
                pos_end=None, 
                use_error_term=self.use_error_term,
                cache=unpatch_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
            
            current_grad = current_grad + grad_through_attn # resid_pre_grad
            
            if self.check_resid_pre:
                # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                bwd_cache[resid_pre_name] = SparseAct(
                    act=current_grad @ self.dict_saes[resid_pre_name].W_dec.T,
                    res = current_grad if self.use_error_term else None,
                )
                
        return bwd_cache 
    
    def _gradient_wrt_edges(
        self,
        clean_token: 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()
        self.model.reset_saes()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()
            
        gradient_mode = kwargs.get('edge_gradient_mode', 'virtual_weight')
        
        # activation cache on clean token
        _, unpatch_clean_cache = self.model.run_with_cache(clean_token)
        # SAE act cache on clean token
        _, clean_cache = self.forward(clean_token, 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 == 'virtual_weight':
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution_trans(
                            unpatch_clean_cache,
                            hook_position_end,
                            hook_position_start,
                            node_grads[hook_position_end],
                            right_vec,
                            **kwargs,
                        )
                    elif gradient_mode == 'gradient':
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution(
                            clean_token,
                            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_trans(
        self,
        unpatched_clean_cache: ActivationCache | Dict[str, Tensor],
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        **kwargs,
    ) -> Tensor:
        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
        
        aggregate_dim = [0] if self.token_wise else [0, 1]
        edge_effect = {}        
        all_error = []
        
        for end_node, end_index in self.active_nodes(*hook_position_end):
            if isinstance(end_index, ErrorIndex):
                all_error.append((end_node, end_index))
            elif isinstance(end_index, FeatureIndex):
                feat_id = end_index.idx[-1]
                pos_end = end_index.idx[-2]
                index = t.tensor(list(end_index.idx), device=self.device)
                end_node_grad = leftvec.act[:, pos_end, feat_id].unsqueeze(-1).unsqueeze(-1) # (b, 1, 1)
                
                end_feature_vec = self.dict_saes[end_node.name].W_enc[:, feat_id].unsqueeze(0).unsqueeze(0) # (1, 1, d_model)
                grad_dot_leftvec_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                    grad=end_feature_vec,
                    pos_end=pos_end,
                    batch_size=leftvec.act.shape[0],
                    unpatched_clean_cache=unpatched_clean_cache,
                    hook_position_end=hook_position_end,
                    hook_position_start=hook_position_start,
                ) * end_node_grad
                grad_dot_leftvec = SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae)
                    act=grad_dot_leftvec_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                    res=grad_dot_leftvec_tensor if self.use_error_term else None
                )
                '''
                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[index] = ( # (seq, d_sae+1) || (d_sae+1)
                    grad_dot_leftvec @ rightvec
                ).sum(aggregate_dim).to_tensor()
            else:
                raise ValueError(f"end_index of type {type(end_index)} is not supported.")
            
        '''
        The gradient of error node to upstream node is 1 - gradient of sum end_feature_node
        We multiply error grad so that we only have to backward once (Jacobian vector product)
        The "end_feature_dependent" sums all of the gradient of end_feature_node.
        '''
        if self.use_error_term:
            all_end_node_grad: Tensor = leftvec.res # type: ignore
            feature_coef_to_cal_error_edge = einops.einsum(
                all_end_node_grad, self.dict_saes[hook_position_end[0].name].W_dec,
                "b seq d_model, d_sae_end d_model -> b seq d_sae_end",
            )
            
            end_feature_dependent = einops.einsum(
                feature_coef_to_cal_error_edge, self.dict_saes[hook_position_end[0].name].W_enc,
                "b seq d_sae_end, d_model d_sae_end -> b seq d_model"
            )
            
            for end_error_node, end_error_index in all_error:
                pos_end = end_error_index.idx[0]
                index = t.tensor(list(end_error_index.idx + (d_sae_end,)), device=self.device)
                end_node_grad = leftvec.res[:, pos_end].unsqueeze(1) # (b, 1, d_model) # type: ignore
                
                end_error_grad_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                    grad=end_node_grad - end_feature_dependent[:, pos_end].unsqueeze(1), # (b, 1, d_model)
                    pos_end=pos_end,
                    batch_size=leftvec.act.shape[0],
                    unpatched_clean_cache=unpatched_clean_cache,
                    hook_position_end=hook_position_end,
                    hook_position_start=hook_position_start,
                )
                end_error_grad = SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae)
                    act=end_error_grad_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                    res=end_error_grad_tensor if self.use_error_term else None
                )

                edge_effect[index] = ( # (seq, d_sae_start+1) | (d_sae_start+1)
                    end_error_grad @ rightvec
                ).sum(aggregate_dim).to_tensor()  
                    
        seq = int(self.seq_length) # 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(list(edge_effect.keys()), dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([value for value 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, 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, num_end, seq, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()
    
    def _edge_attribution(
        self,
        clean_token: Tensor,
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        layer: int,
        **kwargs,
    ) -> Tensor:
        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 = {}
        with t.set_grad_enabled(True):
            with self._detach_error_term(False, hook_position_end[0].name):
                with self._setup_forward_model_hook(transfer_grad=False): # must be false, no gradient from intermediate component
                    with self._setup_fwd_bwd_edge_grad_sae_hook(
                        bwd_cache=bwd_cache,
                        to_bwd_cache=to_bwd_cache,
                        hook_position_start=hook_position_start,
                        hook_position_end=hook_position_end,
                    ):
                        self.model.forward(clean_token, return_type=None, stop_at_layer=layer+1)
                        
                        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 end_node, end_index in 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[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 = int(self.seq_length) # 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(list(edge_effect.keys()), dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([value for value 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, 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, num_end, seq, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()
    
    def _DE_using_virtual_weight(
        self,
        grad: Tensor, # (1, 1, d_model) or (b, 1, d_model)
        pos_end: int,
        batch_size: int,
        unpatched_clean_cache: ActivationCache | Dict[str, Tensor],
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
    ):
        start_layer = int(hook_position_start[0].name.split(".")[1])
        end_layer = int(hook_position_end[0].name.split(".")[1])
        
        path = [] # path attention for gradient, if we have attention circuit then the path is empty
        if not self.check_attn_out:
            if not "resid_pre" in hook_position_end[0].name:
                path.append(("attn", end_layer)) # attn or mlp at end hook -> goes through attn at end_layer
                
            for layer in reversed(range(start_layer+1, end_layer)):
                path.append(("attn", layer))
                
            if "resid_pre" in hook_position_start[0].name:
                path.append(("attn", start_layer)) # resid_pre at start hook -> goes through attn at start_layer
                
        if "mlp_out" in hook_position_end[0].name:
            ln = "ln2"
        elif "attn_out" in hook_position_end[0].name:
            ln = "ln1"
        else:
            ln = None
        
        if self.check_attn_out and "attn_out" in hook_position_end[0].name:
            current_grad = gradient_with_attn(
                model=self.model,
                end_feature_vec=grad, # (b, seq, d_model)
                start_feature_vec=None,
                layer_end=end_layer,
                pos_end=pos_end, 
                use_error_term=self.use_error_term,
                cache=unpatched_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
        else:
            current_grad = gradient_ln_only( # (b, seq, d_model) 
                end_feature_vec=grad,
                start_feature_vec=None,
                layer_end=end_layer,
                pos_end=pos_end,
                seq_length=self.seq_length, # type: ignore
                batch_size=batch_size,
                use_error_term=self.use_error_term,
                cache=unpatched_clean_cache,
                device=self.device, # type: ignore
                ln=ln,
            ).act 
        
        for _, layer in path:
            grad_through_attn = gradient_with_attn(
                model=self.model,
                end_feature_vec=current_grad, # (b, seq, d_model)
                start_feature_vec=None,
                layer_end=layer,
                pos_end=None, 
                use_error_term=self.use_error_term,
                cache=unpatched_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
            
            current_grad = current_grad + grad_through_attn
        
        return current_grad
    
    def run_model(
        self, 
        toks: Tensor, 
        use_error_term: bool | None = None,
    ) -> Tuple[Tensor, ActivationCache]:
        """
        Runs an MLP transcoder(s) on a batch of tokens.
        """
        fwd_cache = {}
        with self._setup_forward_model_hook(use_error_term=use_error_term, transfer_grad=False):
            with self._setup_fwd_sae_hook(
                fwd_cache=fwd_cache, corrupt_cache=None, patch_deleted_comp=False, use_error_term=use_error_term
            ):
                logits = self.model(toks)

        return logits, ActivationCache(cache_dict=fwd_cache, model=self.model) # type: ignore
    
    @contextmanager
    def _setup_fwd_sae_hook(
        self, 
        fwd_cache: Dict[str, Tensor],
        corrupt_cache: ActivationCache | Dict[str, Tensor] | None, 
        patch_deleted_comp: bool = False,
        use_error_term: bool | None = None,
    ):    
        use_error_term = use_error_term if use_error_term is not None else self.use_error_term
        
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == "hook_sae_acts_post" or 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 == "hook_sae_error" or hook.name == error_term_name(sae_name)) and 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
        
        try:
            with self._setup_error_term(use_error_term):
                for sae_name, sae in self.dict_saes.items():
                    sae.add_hook(
                        lambda name: True,
                        partial(hook_sae_fwd, sae_name=sae_name),
                        dir="fwd",
                    )
                yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
                
    @contextmanager
    def _setup_fwd_bwd_grad_sae_hook(
        self, 
        fwd_cache: Dict[str, Tensor],
        bwd_cache: Dict[str, Tensor],
        pass_through_grad: bool,
    ):
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif (hook.name == "hook_sae_error" or error_term_name(sae_name) == hook.name) and self.use_error_term:
                fwd_cache[error_term_name(sae_name)] = act.detach()
                    
            return act
        
        pass_through_cache = {}
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                bwd_cache[sae_hook_name(sae_name)] = grad.detach()
                
            elif hook.name == "hook_sae_output" or 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
                    bwd_cache[error_term_name(sae_name)] = grad.detach()
                # avoid pass through grad for transcoder
                if pass_through_grad and sae_name not in self.output_hooks:
                    pass_through_cache[output_hook_name(sae_name)] = grad.detach()
                    
            elif hook.name == "hook_sae_input" or hook.name == input_hook_name(sae_name):
                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
                    if sae_name not in self.output_hooks:
                        # no pass through grad for transcoder
                        grad.copy_(pass_through_cache[output_hook_name(sae_name)])
                    else:
                        # avoid gradient flow to upstream from the transcoder
                        grad.zero_() 
        
        try:
            with self._setup_error_term(self.use_error_term):
                for sae_name, sae in self.dict_saes.items():
                    sae.add_hook(
                        lambda name: True,
                        partial(hook_sae_fwd, sae_name=sae_name),
                        dir="fwd",
                    )
                    sae.add_hook(
                        lambda name: True,
                        partial(hook_sae_bwd, sae_name=sae_name),
                        dir="bwd",
                    )
                yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
         
    @contextmanager
    def _setup_fwd_bwd_grad_sae_hook_ig(
        self, 
        target_name: str,
        frac: float,
        fwd_cache: Dict[str, Tensor],
        bwd_cache: Dict[str, Tensor],
        corrupt_cache: ActivationCache | Dict[str, Tensor], 
    ):  
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, target_name: str, frac: float) -> Tensor:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                # interpolate for integrated gradients
                if hook.name == sae_hook_name(target_name) or _hook_name(sae_name, 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 (hook.name == "hook_sae_error" or error_term_name(sae_name) == hook.name) and self.use_error_term:
                # interpolate for integrated gradients
                if hook.name == error_term_name(target_name) or _hook_name(sae_name, 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
        
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == "hook_sae_acts_post" or 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 == "hook_sae_output" or 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))
                    bwd_cache[error_term_name(sae_name)] += grad.detach()
        
        try:
            with self._setup_error_term(self.use_error_term):
                for sae_name, sae in self.dict_saes.items():
                    sae.add_hook(
                        lambda name: True,
                        partial(hook_sae_fwd, sae_name=sae_name, target_name=target_name, frac=frac),
                        dir="fwd",
                    )
                    sae.add_hook(
                        lambda name: True,
                        partial(hook_sae_bwd, sae_name=sae_name),
                        dir="bwd",
                    )
                yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
                
    @contextmanager
    def _setup_virtual_weight_sae_hook_ig(
        self, 
        target_name: str,
        frac: float,
        corrupt_cache: ActivationCache | Dict[str, Tensor], 
    ):  
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, target_name: str, frac: float) -> Tensor:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                # interpolate for integrated gradients
                if hook.name == sae_hook_name(target_name) or _hook_name(sae_name, hook.name) == sae_hook_name(target_name):
                    act = interpolate(
                        corrupt_cache[sae_hook_name(sae_name)], 
                        act, 
                        frac,
                    )
                
            elif (hook.name == "hook_sae_error" or error_term_name(sae_name) == hook.name) and self.use_error_term:
                # interpolate for integrated gradients
                if hook.name == error_term_name(target_name) or _hook_name(sae_name, hook.name) == error_term_name(target_name):
                    act = interpolate(
                        corrupt_cache[error_term_name(sae_name)], 
                        act, 
                        frac,
                    )
                    
            return act
        
        try:
            with self._setup_error_term(self.use_error_term):
                for sae_name, sae in self.dict_saes.items():
                    sae.add_hook(
                        lambda name: True,
                        partial(hook_sae_fwd, sae_name=sae_name, target_name=target_name, frac=frac),
                        dir="fwd",
                    )
                yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
              
    @contextmanager
    def _setup_fwd_bwd_edge_grad_sae_hook(
        self, 
        bwd_cache: Dict[str, Tensor],
        to_bwd_cache: Dict[str, Tensor],
        hook_position_start: Tuple[Node, Index],
        hook_position_end: Tuple[Node, Index],
    ):
        # print(hook_position_end[0].name, hook_position_start[0].name)
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str, bwd_cache: Dict) -> None:
            # print(sae_name, hook.name, t.sum(grad != 0).item())
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                if hook.name == sae_hook_name(hook_position_start[0].name) or _hook_name(sae_name, hook.name) == sae_hook_name(hook_position_start[0].name):
                    # only store the gradient at the start position
                    bwd_cache[sae_hook_name(sae_name)] = grad.detach()
                
            elif hook.name == "hook_sae_output" or hook.name == output_hook_name(sae_name):
                if hook.name == output_hook_name(hook_position_start[0].name) or _hook_name(sae_name, 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
                        bwd_cache[error_term_name(sae_name)] = grad.detach()

            # 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 == "hook_sae_input" or hook.name == input_hook_name(sae_name):
                if hook.name == input_hook_name(hook_position_end[0].name) or _hook_name(sae_name, hook.name) == input_hook_name(hook_position_end[0].name):
                    pass # no zero grad for end position
                else:
                    grad.zero_()

        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, to_bwd_cache: Dict) -> Tensor:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                if  hook.name == sae_hook_name(hook_position_end[0].name) or _hook_name(sae_name, hook.name) == sae_hook_name(hook_position_end[0].name):
                    # store activation for backward later
                    to_bwd_cache[sae_hook_name(sae_name)] = act
                
            elif (hook.name == "hook_sae_error" or error_term_name(sae_name) == hook.name) and self.use_error_term:
                if  hook.name == error_term_name(hook_position_end[0].name) or _hook_name(sae_name, hook.name) == error_term_name(hook_position_end[0].name):
                    # store activation for backward later
                    to_bwd_cache[error_term_name(sae_name)] = act
                    
            return act
        
        try:
            with self._setup_error_term(self.use_error_term):
                for sae_name, sae in self.dict_saes.items():
                    sae.add_hook(
                        lambda name: True,
                        partial(hook_sae_fwd, sae_name=sae_name, to_bwd_cache=to_bwd_cache),
                        dir="fwd",
                    )
                    sae.add_hook(
                        lambda name: True,
                        partial(hook_sae_bwd, sae_name=sae_name, bwd_cache=bwd_cache),
                        dir="bwd",
                    )
                yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
                
    @contextmanager
    def _setup_forward_model_hook(self, use_error_term: bool | None = None, transfer_grad: bool = True):
        trans_cache = {}
        # Hook function at transcoder input: caches the activations before transcoder
        def hook_transcoder_input(activations: Tensor, hook: HookPoint, transcoder_idx: int):
            trans_cache[transcoder_idx] = activations.clone()

        # Hook function at transcoder output: replaces activations with transcoder output
        def hook_transcoder_output(activations: Tensor, hook: HookPoint, transcoder_idx: int):
            trans_input = trans_cache[transcoder_idx]
            if transfer_grad:
                # allow gradient to pass through the transcoder output and original activations
                return self.transcoders[transcoder_idx]((trans_input, activations)) + (activations - activations.detach())
            else:
                # no original activations gradient, but save memory
                return self.transcoders[transcoder_idx]((trans_input, activations))
        
        fwd_hooks = []
        for i in range(len(self.transcoders)):
            fwd_hooks.append((self.input_hooks[i], partial(hook_transcoder_input, transcoder_idx=i)))
            fwd_hooks.append((self.output_hooks[i], partial(hook_transcoder_output, transcoder_idx=i)))
            
        use_error_term = use_error_term if use_error_term is not None else self.use_error_term
        try: 
            for sae in self._saes_to_list():
                self.model.add_sae(sae, use_error_term)
            for hook, func in fwd_hooks:
                self.model.add_hook(hook, func, dir="fwd")
            
            yield
        finally:
            self.model.reset_saes()
            self.model.reset_hooks()
    
    @contextmanager  
    def _setup_error_term(self, use_error_term: bool | None = None):
        if use_error_term is None:
            use_error_term = self.use_error_term
        
        orig_use_error_term = {}
        try: 
            for sae_name, sae in self.dict_saes.items():
                sae.use_error_term = use_error_term
                orig_use_error_term[sae_name] = sae.use_error_term
            yield
        finally:
            for sae_name, sae in self.dict_saes.items():
                sae.use_error_term = orig_use_error_term[sae_name]   
            
    def _saes_to_list(self) -> List[Any]:
        return self.non_transcoders
    

class ESAE_FG_Trans(
    ESAE_FG,
    Feature_Graph_Trans,
):
    def __init__(
        self,
        model: HookedSAETransformer,
        saes: Dict[int, List[Tuple[str, Any]]],
        esaes: Dict[int, List[Tuple[str, Any]]],
        use_esae_error_term: bool = False,
    ) -> None:
        
        super().__init__(model, saes, esaes, use_esae_error_term)
        
        self.process_transcoder()
                
    def forward(
        self,
        clean_token: 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 = {}
        
        with t.no_grad():
            with self._setup_forward_model_hook(transfer_grad=False):
                with self._setup_fwd_sae_hook(
                    fwd_cache=fwd_cache, corrupt_cache=corrupt_cache, patch_deleted_comp=patch_deleted_comp
                ):
                    logits = self.model(clean_token)
        
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                fwd_cache,
                sae_hook_name(sae_name),
                sae_hook_name(error_term_name(sae_name)) if self.use_error_term else None,
                error_term_name(error_term_name(sae_name)) if self.use_esae_error_term else None,
            )
            
        for sae in self.dict_saes.values():
            sae.reset_hooks()
        self.model.reset_hooks()
        
        return logits, cache
    
    def forward_backward_gradient(
        self,
        clean_token: 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
    ]:
        return Feature_Graph_Trans.forward_backward_gradient(
            self,
            clean_token=clean_token,
            corrupt_cache=corrupt_cache,
            metric=metric,
            retain_graph=retain_graph,
            mode=mode,
            gradient_mode=gradient_mode,
            pass_through_grad=pass_through_grad,
            verbose=verbose,
            **kwargs,
        )
    
    def _gradient_wrt_nodes(
        self,
        clean_token: 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]
    ]:
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()
            
        fwd_cache = {}
        bwd_cache = {}
            
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self._setup_forward_model_hook(transfer_grad=kwargs.get("transfer_grad", True)):
                    with self._setup_fwd_bwd_grad_sae_hook(
                        fwd_cache=fwd_cache, 
                        bwd_cache=bwd_cache, 
                        pass_through_grad=pass_through_grad,
                    ):
                        metric(self.model(clean_token)).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),
                sae_hook_name(error_term_name(node.name)) if self.use_error_term else None,
                error_term_name(error_term_name(node.name)) if self.use_esae_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),
                sae_hook_name(error_term_name(sae_name)) if self.use_error_term else None,
                error_term_name(error_term_name(sae_name)) if self.use_esae_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,
        clean_token: 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()

        fwd_cache = {}
        bwd_cache = {}
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self._setup_forward_model_hook(transfer_grad=kwargs.get("transfer_grad", True)):
                    for target_name in self.dict_saes.keys():
                        for step in range(steps):
                            frac = step / steps
                            with self._setup_fwd_bwd_grad_sae_hook_ig(
                                target_name=target_name,
                                frac=frac,
                                fwd_cache=fwd_cache,
                                bwd_cache=bwd_cache,
                                corrupt_cache=corrupt_cache,
                            ):
                                metric(self.model(clean_token)).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),
                sae_hook_name(error_term_name(node.name)) if self.use_error_term else None,
                error_term_name(error_term_name(node.name)) if self.use_esae_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),
                sae_hook_name(error_term_name(sae_name)) if self.use_error_term else None,
                error_term_name(error_term_name(sae_name)) if self.use_esae_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_vw(
        self,
        clean_token: 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]
    ]:
        '''
        Using virtual weight to compute node grad
        '''
        self._check_graph()
        self.model_setup()
        
        self.model.reset_hooks()
            
        # unpatched cache on clean token
        _, unpatch_clean_cache = self.model.run_with_cache(clean_token)
        
        sink_hook_name = get_act_name("resid_post", self.n_layers-1) # last layer resid post hook
        
        sink_node_cache: Dict[str, Tensor] = {}
        def hook_bwd(tens: Tensor, hook: HookPoint):
            sink_node_cache[hook.name] = tens.detach() # type: ignore
            
            
        fwd_cache = {}
        with t.set_grad_enabled(True):
            with self._detach_error_term(True):
                with self.model.hooks(
                    bwd_hooks=[(sink_hook_name, hook_bwd)]
                ):
                    with self._setup_forward_model_hook(transfer_grad=False):
                        with self._setup_fwd_bwd_grad_sae_hook(
                            fwd_cache=fwd_cache, 
                            bwd_cache={}, 
                            pass_through_grad=pass_through_grad,
                        ):
                            metric(self.model(clean_token)).backward(retain_graph=retain_graph)
                        
        current_grad = sink_node_cache[sink_hook_name] # (b, seq, d_model)
        bwd_cache = {}
        for layer in reversed(range(self.n_layers)):
            mlp_name = get_act_name("mlp_out", layer)
            attn_name = get_act_name("attn_out", layer)
            resid_pre_name = get_act_name("resid_pre", layer)
            
            # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
            bwd_cache[mlp_name] = SparseAct(
                act=current_grad @ self.dict_saes[mlp_name].W_dec.T,
                res = current_grad @ self.dict_esaes[mlp_name].W_dec.T if self.use_error_term else None,
                resc = current_grad if self.use_esae_error_term else None,
            )
            
            grad_through_mlp = gradient_with_mlp(
                end_feature_vec=current_grad, # (b, seq, d_model)
                start_feature_vec=None,
                transcoder_enc=self.dict_saes[mlp_name].W_enc, # (d_model, d_sae)
                transcoder_dec=self.dict_saes[mlp_name].W_dec, # (d_sae, d_model)
                transcoder_act=fwd_cache[sae_hook_name(mlp_name)], # (b, seq, d_sae)
                layer_end=layer,
                pos_end=None, 
                seq_length=self.seq_length, # type: ignore
                batch_size=current_grad.shape[0],
                use_error_term=self.use_error_term,
                cache=unpatch_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
            
            current_grad = current_grad + grad_through_mlp # resid_mid_grad
            
            if self.check_attn_out:
                # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                bwd_cache[attn_name] = SparseAct(
                    act=current_grad @ self.dict_saes[attn_name].W_dec.T,
                    res = current_grad @ self.dict_esaes[attn_name].W_dec.T if self.use_error_term else None,
                    resc = current_grad if self.use_esae_error_term else None,
                )
            
            grad_through_attn = gradient_with_attn(
                model=self.model,
                end_feature_vec=current_grad, # (b, seq, d_model)
                start_feature_vec=None,
                layer_end=layer,
                pos_end=None, 
                use_error_term=self.use_error_term,
                cache=unpatch_clean_cache,
                device=self.device # type: ignore
            ).act # (b, seq, d_model)
            
            current_grad = current_grad + grad_through_attn # resid_pre_grad
            
            if self.check_resid_pre:
                # (b, seq, d_model) @ (d_model, d_sae) -> (b, seq, d_sae)
                bwd_cache[resid_pre_name] = SparseAct(
                    act=current_grad @ self.dict_saes[resid_pre_name].W_dec.T,
                    res = current_grad @ self.dict_esaes[resid_pre_name].W_dec.T if self.use_error_term else None,
                    resc = current_grad if self.use_esae_error_term else None,
                )
            
        node_grads = {}
        for node, index in self.nodes.keys():
            node_grads[(node, index)] = bwd_cache[node.name]
            
        cache = {}
        for sae_name in self.dict_saes.keys():
            cache[sae_name] = cache_to_sparseact(
                fwd_cache,
                sae_hook_name(sae_name),
                sae_hook_name(error_term_name(sae_name)) if self.use_error_term else None,
                error_term_name(error_term_name(sae_name)) if self.use_esae_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,
        clean_token: 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()
        self.model.reset_saes()
        for _, sae in self.dict_saes.items():
            sae.reset_hooks()
            
        gradient_mode = kwargs.get('edge_gradient_mode', 'virtual_weight')
        
        # unpatched cache on clean token
        _, unpatch_clean_cache = self.model.run_with_cache(clean_token)
        # potentially patched cache on clean token
        _, clean_cache = self.forward(clean_token, 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), 
                        sae_hook_name(error_term_name(hook_position_start[0].name)) if self.use_error_term else None,
                        error_term_name(error_term_name(hook_position_start[0].name)) if self.use_esae_error_term else None,
                    )
                    right_vec = corrupt_sparse_act - clean_cache[hook_position_start[0].name]
                    
                    if gradient_mode == "virtual_weight":
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution_trans(
                            unpatch_clean_cache,
                            hook_position_end,
                            hook_position_start,
                            node_grads[hook_position_end],
                            right_vec,
                            **kwargs,
                        )
                    elif gradient_mode == 'gradient':
                        edge_grads[hook_position_start + hook_position_end] = self._edge_attribution(
                            clean_token,
                            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_trans(
        self,
        unpatched_clean_cache: ActivationCache | Dict[str, Tensor],
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        **kwargs,
    ) -> Tensor:
        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
        d_esae_end = self.dict_esaes[hook_position_end[0].name].cfg.d_sae
        d_esae_start = self.dict_esaes[hook_position_start[0].name].cfg.d_sae
        
        aggregate_dim = [0] if self.token_wise else [0, 1]
        edge_effect = {}        
        all_error = []
        all_feature_error = []
        
        for end_node, end_index in self.active_nodes(*hook_position_end):
            if isinstance(end_index, ErrorIndex):
                all_error.append((end_node, end_index))
            elif isinstance(end_index, FeatureErrorIndex):
                all_feature_error.append((end_node, end_index))
            elif isinstance(end_index, FeatureIndex):
                feat_id = end_index.idx[-1]
                pos_end = end_index.idx[-2]
                index = t.tensor(list(end_index.idx), device=self.device)
                end_node_grad = leftvec.act[:, pos_end, feat_id].unsqueeze(-1).unsqueeze(-1) # (b, 1, 1)
                
                end_feature_vec = self.dict_saes[end_node.name].W_enc[:, feat_id].unsqueeze(0).unsqueeze(0) # (1, 1, d_model)
                grad_dot_leftvec_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                    grad=end_feature_vec,
                    pos_end=pos_end,
                    batch_size=leftvec.act.shape[0],
                    unpatched_clean_cache=unpatched_clean_cache,
                    hook_position_end=hook_position_end,
                    hook_position_start=hook_position_start,
                ) * end_node_grad
                grad_dot_leftvec = SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae)
                    act=grad_dot_leftvec_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                    res=grad_dot_leftvec_tensor @ self.dict_esaes[hook_position_start[0].name].W_dec.T if self.use_error_term else None,
                    resc=grad_dot_leftvec_tensor if self.use_esae_error_term else None,
                )
                '''
                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)
                '''
                effect = (
                    grad_dot_leftvec * rightvec
                ).sum(aggregate_dim)
                if self.use_esae_error_term:
                    effect.contract()
                edge_effect[index] = effect.to_tensor() # (seq, d_sae_start+d_esae_start+1) | (d_sae_start+d_esae_start+1)
            else:
                raise ValueError(f"end_index of type {type(end_index)} is not supported.")
            
        '''
        The gradient of feature error node to upstream node is f_esae_enc - sum gradient of end_feature_node
        The "grad_through_end_feat_error" computes the jacobian of f_esae_enc going through f_sae_dec and f_sae_enc
        '''
        if self.use_error_term:            
            grad_through_end_feat_error = einops.einsum( 
                self.dict_saes[hook_position_end[0].name].W_dec.T, self.dict_saes[hook_position_end[0].name].W_enc.T,
                "d_model1 d_sae_end, d_sae_end d_model2 -> d_model1 d_model2"
            )
            grad_through_end_feat_error = einops.einsum( # gradient through sae_feature
                self.dict_esaes[hook_position_end[0].name].W_enc.T, grad_through_end_feat_error,
                "d_esae_end d_model1, d_model1 d_model2 -> d_esae_end d_model2"
            )
            
            for end_feature_error_node, end_feature_error_index in all_feature_error:
                pos_end = end_feature_error_index.idx[0]
                feat_id = end_feature_error_index.idx[-1]
                
                revised_index = list(end_feature_error_index.idx)
                revised_index[-1] += d_sae_end
                index = t.tensor(revised_index, device=self.device)
                
                end_node_grad = leftvec.res[:, pos_end, feat_id].unsqueeze(-1).unsqueeze(-1) # (b, 1, 1) # type: ignore
                end_feature_vec = self.dict_esaes[end_feature_error_node.name].W_enc[:, feat_id].unsqueeze(0).unsqueeze(0) # (1, 1, d_model)
                
                end_feature_error_grad_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                    grad=end_feature_vec - grad_through_end_feat_error[feat_id].unsqueeze(0).unsqueeze(0), # (1, 1, d_model)
                    pos_end=pos_end,
                    batch_size=leftvec.act.shape[0],
                    unpatched_clean_cache=unpatched_clean_cache,
                    hook_position_end=hook_position_end,
                    hook_position_start=hook_position_start,
                )
                end_feature_error_grad = SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae)
                    act=end_feature_error_grad_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                    res=end_feature_error_grad_tensor @ self.dict_esaes[hook_position_start[0].name].W_dec.T,
                    resc=end_feature_error_grad_tensor if self.use_esae_error_term else None,
                )
                effect = (
                    end_feature_error_grad * rightvec
                ).sum(aggregate_dim)
                if self.use_esae_error_term:
                    effect.contract()
                edge_effect[index] = effect.to_tensor() # (seq, d_sae_start+d_esae_start+1) | (d_sae_start+d_esae_start+1)
                
        '''
        The gradient of feature error node to upstream node is:
            f_esae_enc - sum gradient of end_feature_node - sum gradient of end_feaeture_error_node
            
        
        The gradient of feature error node to upstream node is f_esae_enc - sum gradient of end_feature_node (see above)
        The "grad_through_end_feat_error" computes the jacobian of f_esae_enc going through f_sae_dec and f_sae_enc
        The "feature_error_coef_to_cal_error_edge" calculates the leftvec (metric gradient) at the esae_error | esae_feature
        We can then have the contribution of feature_error_node by: 
            The contribution of feature_end_node is the gradient of f_esae_enc - sum grad_through_end_feat_error * feature_error_coef_to_cal_error_edge
        
        The "end_feature_dependent" sums all of the gradient of end_feature_node.
        '''
                
        if self.use_esae_error_term and self.use_error_term:
            all_end_node_grad: Tensor = leftvec.resc # type: ignore
            '''
            End feauture contribution
            '''
            feature_coef_to_cal_error_edge = einops.einsum(
                all_end_node_grad, self.dict_saes[hook_position_end[0].name].W_dec,
                "b seq d_model, d_sae_end d_model -> b seq d_sae_end",
            )
            
            end_feature_dependent = einops.einsum(
                feature_coef_to_cal_error_edge, self.dict_saes[hook_position_end[0].name].W_enc,
                "b seq d_sae_end, d_model d_sae_end -> b seq d_model"
            )
            
            '''
            End feature error contribution
            '''
            grad_through_end_feat_error = einops.einsum( 
                self.dict_saes[hook_position_end[0].name].W_dec.T, self.dict_saes[hook_position_end[0].name].W_enc.T,
                "d_model1 d_sae_end, d_sae_end d_model2 -> d_model1 d_model2"
            )
            grad_through_end_feat_error = einops.einsum( # gradient through sae_feature
                self.dict_esaes[hook_position_end[0].name].W_enc.T, grad_through_end_feat_error,
                "d_esae_end d_model1, d_model1 d_model2 -> d_esae_end d_model2"
            )
            
            feature_error_coef_to_cal_error_edge = einops.einsum(
                all_end_node_grad, self.dict_esaes[hook_position_end[0].name].W_dec,
                "b seq d_model, d_esae_end d_model -> b seq d_esae_end",
            )
            
            feature_error_dependent = einops.einsum(
                self.dict_esaes[hook_position_end[0].name].W_enc.T - grad_through_end_feat_error,
                feature_error_coef_to_cal_error_edge,
                "d_esae_end d_model, b seq d_esae_end -> b seq d_model",
            )
            
            for end_error_node, end_error_index in all_error:
                pos_end = end_error_index.idx[0]
                index = t.tensor(list(end_error_index.idx + (d_sae_end+d_esae_end,)), device=self.device)
                end_node_grad = leftvec.resc[:, pos_end].unsqueeze(1) # (b, 1, d_model) # type: ignore
                
                end_error_grad_tensor = self._DE_using_virtual_weight( # (b, seq, d_model)
                    # (b, 1, d_model)
                    grad= end_node_grad - end_feature_dependent[:, pos_end].unsqueeze(1) - feature_error_dependent[:, pos_end].unsqueeze(1), 
                    pos_end=pos_end,
                    batch_size=leftvec.act.shape[0],
                    unpatched_clean_cache=unpatched_clean_cache,
                    hook_position_end=hook_position_end,
                    hook_position_start=hook_position_start,
                )
                end_error_grad = SparseAct(
                    # (b, seq, d_model) @ (d_model, d_sae)
                    act=end_error_grad_tensor @ self.dict_saes[hook_position_start[0].name].W_dec.T,
                    res=end_error_grad_tensor @ self.dict_esaes[hook_position_start[0].name].W_dec.T,
                    resc=end_error_grad_tensor,
                )

                effect = (
                    end_error_grad * rightvec
                ).sum(aggregate_dim)
                if self.use_esae_error_term:
                    effect.contract()
                edge_effect[index] = effect.to_tensor() # (seq, d_sae_start+d_esae_start+1) | (d_sae_start+d_esae_start+1)
                    
        seq = int(self.seq_length) # type: ignore
        num_end = d_sae_end
        num_start = d_sae_start
        if self.use_error_term:
            num_end += d_esae_end
            num_start += d_esae_start
        if self.use_esae_error_term:
            num_end += 1
            num_start += 1
        
        if len(edge_effect.keys()) != 0:
            indices = t.stack(list(edge_effect.keys()), dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([value for value in edge_effect.values()], dim=0) # shape (num_active, seq, d_sae+d_esae+1) or (num_active, d_sae+d_esae+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, 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, num_end, seq, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()
    
    def _edge_attribution(
        self,
        clean_token: Tensor,
        hook_position_end: Tuple[Node, Index],
        hook_position_start: Tuple[Node, Index],
        leftvec: SparseAct,
        rightvec: SparseAct,
        layer: int,
        **kwargs,
    ) -> Tensor:
        d_sae_end = self.dict_saes[hook_position_end[0].name].cfg.d_sae
        d_esae_end = self.dict_esaes[hook_position_end[0].name].cfg.d_sae
        d_sae_start = self.dict_saes[hook_position_start[0].name].cfg.d_sae
        d_esae_start = self.dict_esaes[hook_position_start[0].name].cfg.d_sae
        
        to_bwd_cache = {}
        bwd_cache = {}
        edge_effect = {}
        with t.set_grad_enabled(True):
            with self._detach_error_term(False, hook_position_end[0].name):
                with self._setup_forward_model_hook(transfer_grad=False): # must be false, no gradient from intermediate component
                    with self._setup_fwd_bwd_edge_grad_sae_hook(
                        bwd_cache=bwd_cache,
                        to_bwd_cache=to_bwd_cache,
                        hook_position_start=hook_position_start,
                        hook_position_end=hook_position_end,
                    ):
                        self.model.forward(clean_token, return_type=None, stop_at_layer=layer+1)
                        
                        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), 
                                sae_hook_name(error_term_name(hook_position_end[0].name)) if self.use_error_term else None,
                                error_term_name(error_term_name(hook_position_end[0].name)) if self.use_esae_error_term else None,
                            ) * leftvec.detach()
                        ).sum(aggregate_dim)
                        if self.use_esae_error_term:
                            to_bwd = to_bwd.contract()
                        to_bwd = to_bwd.to_tensor() # (seq, d_sae+d_esae+1) || (d_sae+d_esae+1)
                        
                        del to_bwd_cache
                        
                        for end_node, end_index in self.active_nodes(*hook_position_end):
                            if isinstance(end_index, ErrorIndex):
                                # the last index is error: shape (d_sae+d_esae+1) so last index is d_sae+d_esae
                                to_bwd[end_index.idx + (d_sae_end+d_esae_end,)].backward(retain_graph=True)
                                index = t.tensor(list(end_index.idx + (d_sae_end+d_esae_end,)), device=self.device)
                            elif isinstance(end_index, FeatureErrorIndex):
                                revised_index = list(end_index.idx)
                                revised_index[-1] += d_sae_end
                                to_bwd[tuple(revised_index)].backward(retain_graph=True)
                                index = t.tensor(revised_index, 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+d_esae+1, seq, d_sae+d_esae+1) or (d_sae+d_esae+1, d_sae+d_esae+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+d_esae+1) or (num_active, d_sae+d_esae+1)
                            '''
                            effect = ( 
                                cache_to_sparseact(
                                    bwd_cache, 
                                    sae_hook_name(hook_position_start[0].name), 
                                    sae_hook_name(error_term_name(hook_position_start[0].name)) if self.use_error_term else None,
                                    error_term_name(error_term_name(hook_position_start[0].name)) if self.use_esae_error_term else None,
                                ) * rightvec
                            ).sum(aggregate_dim)
                            if self.use_esae_error_term:
                                effect.contract()
                            edge_effect[index] = effect.to_tensor() # (seq, d_sae+d_esae+1) || (d_sae+d_esae+1)
                        
                        del bwd_cache
                    
        seq = int(self.seq_length) # type: ignore
        num_end = d_sae_end
        num_start = d_sae_start
        if self.use_error_term:
            num_end += d_esae_end
            num_start += d_esae_start
        if self.use_esae_error_term:
            num_end += 1
            num_start += 1
        
        if len(edge_effect.keys()) != 0:
            indices = t.stack(list(edge_effect.keys()), dim=0).T # shape (2, num_active) or (1, num_active)
            values = t.stack([value for value in edge_effect.values()], dim=0) # shape (num_active, seq, d_sae+d_esae+1) or (num_active, d_sae+d_esae+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, 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, num_end, seq, num_start)).coalesce()
        else:
            return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce()
    
    @contextmanager
    def _setup_fwd_sae_hook(
        self, 
        fwd_cache: Dict[str, Tensor],
        corrupt_cache: ActivationCache | Dict[str, Tensor] | None, 
        patch_deleted_comp: bool = False,
        use_error_term: bool | None = None,
    ):    
        use_esae_error_term = use_error_term if use_error_term is not None else self.use_esae_error_term
        
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == "hook_sae_acts_post" or 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 == "hook_sae_error.hook_sae_acts_post" or hook.name == sae_hook_name(error_term_name(sae_name))
            ):
                if patch_deleted_comp and corrupt_cache is not None:
                    act_mask = self.nodes[(sae_name, (None,))].res == 0 # type: ignore
                    act[:, act_mask] = corrupt_cache[sae_hook_name(error_term_name(sae_name))][:, act_mask]
                fwd_cache[sae_hook_name(error_term_name(sae_name))] = act.detach()
                
            elif use_esae_error_term and (
                hook.name == "hook_sae_error.hook_sae_error" or hook.name == error_term_name(error_term_name(sae_name))
            ):
                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(error_term_name(sae_name))][:, resc_mask]
                fwd_cache[error_term_name(error_term_name(sae_name))] = act.detach()
                
            return act
        
        try:
            with self._hook_esaes_to_saes(use_esae_error_term=use_esae_error_term):
                with self._setup_error_term(True):
                    for sae_name, sae in self.dict_saes.items():
                        sae.add_hook(
                            lambda name: True,
                            partial(hook_sae_fwd, sae_name=sae_name),
                            dir="fwd",
                        )
                    yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
                
    @contextmanager
    def _setup_fwd_bwd_grad_sae_hook(
        self, 
        fwd_cache: Dict[str, Tensor],
        bwd_cache: Dict[str, Tensor],
        pass_through_grad: bool,
    ):
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                fwd_cache[sae_hook_name(sae_name)] = act.detach()
                
            elif self.use_error_term and (
                hook.name == "hook_sae_error.hook_sae_acts_post" or hook.name == sae_hook_name(error_term_name(sae_name))
            ):
                fwd_cache[sae_hook_name(error_term_name(sae_name))] = act.detach()
                
            elif self.use_esae_error_term and (
                hook.name == "hook_sae_error.hook_sae_error" or hook.name == error_term_name(error_term_name(sae_name))
            ):
                fwd_cache[error_term_name(error_term_name(sae_name))] = act.detach()
                
            return act
        
        pass_through_cache = {}
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                bwd_cache[sae_hook_name(sae_name)] = grad.detach()
                
            elif hook.name == "hook_sae_output" or hook.name == output_hook_name(sae_name):
                # avoid pass through grad for transcoder
                if pass_through_grad and sae_name not in self.output_hooks:
                    pass_through_cache[output_hook_name(sae_name)] = grad.detach()
                    
            elif self.use_error_term and (
                hook.name == "hook_sae_error.hook_sae_acts_post" or hook.name == sae_hook_name(error_term_name(sae_name))
            ):
                bwd_cache[sae_hook_name(error_term_name(sae_name))] = grad.detach()
                
            elif self.use_esae_error_term and (
                hook.name == "hook_sae_error.hook_sae_error" or hook.name == error_term_name(error_term_name(sae_name))
            ):
                # the error_grad is the output_grad, due to: output = recon + error_term
                bwd_cache[error_term_name(error_term_name(sae_name))] = grad.detach()
                    
            elif hook.name == "hook_sae_input" or hook.name == input_hook_name(sae_name):
                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
                    if sae_name not in self.output_hooks:
                        # no pass through grad for transcoder
                        grad.copy_(pass_through_cache[output_hook_name(sae_name)])
                    else:
                        # avoid gradient flow to upstream from the transcoder
                        grad.zero_() 
        
        try:
            with self._hook_esaes_to_saes(use_esae_error_term=self.use_esae_error_term):
                with self._setup_error_term(self.use_error_term):
                    for sae_name, sae in self.dict_saes.items():
                        sae.add_hook(
                            lambda name: True,
                            partial(hook_sae_fwd, sae_name=sae_name),
                            dir="fwd",
                        )
                        sae.add_hook(
                            lambda name: True,
                            partial(hook_sae_bwd, sae_name=sae_name),
                            dir="bwd",
                        )
                    yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
             
    @contextmanager
    def _setup_fwd_bwd_edge_grad_sae_hook(
        self, 
        bwd_cache: Dict[str, Tensor],
        to_bwd_cache: Dict[str, Tensor],
        hook_position_start: Tuple[Node, Index],
        hook_position_end: Tuple[Node, Index],
    ):
        # print(hook_position_end[0].name, hook_position_start[0].name)
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str, bwd_cache: Dict) -> None:
            # print(sae_name, hook.name, t.sum(grad != 0).item())
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                if hook.name == sae_hook_name(hook_position_start[0].name) or _hook_name(sae_name, hook.name) == sae_hook_name(hook_position_start[0].name):
                    # only store the gradient at the start position
                    bwd_cache[sae_hook_name(sae_name)] = grad.detach()
                    
            elif self.use_error_term and (
                hook.name == "hook_sae_error.hook_sae_acts_post" or hook.name == sae_hook_name(error_term_name(sae_name))
            ):
                if hook.name == sae_hook_name(error_term_name(hook_position_start[0].name)) or _hook_name(sae_name, hook.name) == sae_hook_name(error_term_name(hook_position_start[0].name)):
                    # only store the gradient at the start position
                    bwd_cache[sae_hook_name(error_term_name(sae_name))] = grad.detach()
                
            elif self.use_esae_error_term and (
                hook.name == "hook_sae_error.hook_sae_output" or hook.name == output_hook_name(error_term_name(sae_name))
            ):
                if hook.name == output_hook_name(error_term_name(hook_position_start[0].name)) or _hook_name(sae_name, hook.name) == output_hook_name(error_term_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
                    # only store the gradient at the start position
                    bwd_cache[error_term_name(error_term_name(sae_name))] = grad.detach()

            # 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 == "hook_sae_input" or hook.name == input_hook_name(sae_name):
                if hook.name == input_hook_name(hook_position_end[0].name) or _hook_name(sae_name, hook.name) == input_hook_name(hook_position_end[0].name):
                    pass # no zero grad for end position
                else:
                    grad.zero_()

        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, to_bwd_cache: Dict) -> Tensor:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                if  hook.name == sae_hook_name(hook_position_end[0].name) or _hook_name(sae_name, hook.name) == sae_hook_name(hook_position_end[0].name):
                    # store activation for backward later
                    to_bwd_cache[sae_hook_name(sae_name)] = act
                
            elif self.use_error_term and (
                hook.name == "hook_sae_error.hook_sae_acts_post" or sae_hook_name(error_term_name(sae_name)) == hook.name
            ):
                # if not at the end position, we intervene as usual
                if  hook.name == sae_hook_name(error_term_name(hook_position_end[0].name)) or _hook_name(sae_name, hook.name) == sae_hook_name(error_term_name(hook_position_end[0].name)):
                    to_bwd_cache[sae_hook_name(error_term_name(sae_name))] = act
                        
            elif self.use_esae_error_term and (
                hook.name == "hook_sae_error.hook_sae_error" or error_term_name(error_term_name(sae_name)) == hook.name
            ):
                # if not at the end position, we intervene as usual
                if  hook.name == error_term_name(error_term_name(hook_position_end[0].name)) or _hook_name(sae_name, hook.name) == error_term_name(error_term_name(hook_position_end[0].name)):
                    to_bwd_cache[error_term_name(error_term_name(sae_name))] = act
                    
            return act
        
        try:
            with self._hook_esaes_to_saes(use_esae_error_term=self.use_esae_error_term):
                with self._setup_error_term(self.use_error_term):
                    for sae_name, sae in self.dict_saes.items():
                        sae.add_hook(
                            lambda name: True,
                            partial(hook_sae_fwd, sae_name=sae_name, to_bwd_cache=to_bwd_cache),
                            dir="fwd",
                        )
                        sae.add_hook(
                            lambda name: True,
                            partial(hook_sae_bwd, sae_name=sae_name, bwd_cache=bwd_cache),
                            dir="bwd",
                        )
                    yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
                
    @contextmanager
    def _setup_fwd_bwd_grad_sae_hook_ig(
        self, 
        target_name: str,
        frac: float,
        fwd_cache: Dict[str, Tensor],
        bwd_cache: Dict[str, Tensor],
        corrupt_cache: ActivationCache | Dict[str, Tensor], 
    ):  
        def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, target_name: str, frac: float) -> Tensor:
            if hook.name == "hook_sae_acts_post" or hook.name == sae_hook_name(sae_name):
                # interpolate for integrated gradients
                if hook.name == sae_hook_name(target_name) or _hook_name(sae_name, 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 self.use_error_term and (
                hook.name == "hook_sae_error.hook_sae_acts_post" or hook.name == sae_hook_name(error_term_name(sae_name))
            ):
                # interpolate for integrated gradients
                if (
                    hook.name == sae_hook_name(error_term_name(target_name)) or 
                    _hook_name(sae_name, hook.name) == sae_hook_name(error_term_name(target_name))
                ):
                    act = interpolate(
                        corrupt_cache[sae_hook_name(error_term_name(sae_name))], 
                        act, 
                        frac,
                    )
                fwd_cache[sae_hook_name(error_term_name(sae_name))] = act.detach()
                
            elif self.use_esae_error_term and (
                hook.name == "hook_sae_error.hook_sae_error" or hook.name == error_term_name(error_term_name(sae_name))
            ):
                if (
                    hook.name == error_term_name(error_term_name(target_name)) or 
                    _hook_name(sae_name, hook.name) == error_term_name(error_term_name(target_name))
                ):
                    act = interpolate(
                        corrupt_cache[error_term_name(error_term_name(sae_name))], 
                        act, 
                        frac,
                    )
                fwd_cache[error_term_name(error_term_name(sae_name))] = act.detach()
                    
            return act
        
        def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None:
            if hook.name == "hook_sae_acts_post" or 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 self.use_error_term and (
                hook.name == "hook_sae_error.hook_sae_acts_post" or hook.name == sae_hook_name(error_term_name(sae_name))
            ):
                create_list(bwd_cache, sae_hook_name(error_term_name(sae_name)))
                bwd_cache[sae_hook_name(error_term_name(sae_name))] += grad.detach()
                
            elif self.use_esae_error_term and (
                hook.name == "hook_sae_error.hook_sae_error" or hook.name == error_term_name(error_term_name(sae_name))
            ):
                create_list(bwd_cache, error_term_name(error_term_name(sae_name)))
                bwd_cache[error_term_name(error_term_name(sae_name))] += grad.detach()
        
        try:
            with self._hook_esaes_to_saes(use_esae_error_term=self.use_esae_error_term):
                with self._setup_error_term(self.use_error_term):
                    for sae_name, sae in self.dict_saes.items():
                        sae.add_hook(
                            lambda name: True,
                            partial(hook_sae_fwd, sae_name=sae_name, target_name=target_name, frac=frac),
                            dir="fwd",
                        )
                        sae.add_hook(
                            lambda name: True,
                            partial(hook_sae_bwd, sae_name=sae_name),
                            dir="bwd",
                        )
                    yield
        finally:
            for sae in self.dict_saes.values():
                sae.reset_hooks()
                
    def run_model(
        self, 
        toks: Tensor, 
        use_error_term: bool | None = None,
    ) -> Tuple[Tensor, ActivationCache]:
        """
        Runs an MLP transcoder(s) on a batch of tokens.
        """
        return Feature_Graph_Trans.run_model(
            self,
            toks=toks,
            use_error_term=use_error_term
        )
                
                
    def _saes_to_list(self) -> List[Any]:
        return Feature_Graph_Trans._saes_to_list(self)