update

skeleton_modeling_albert.py

@@ -21,27 +21,24 @@ def SkeletonAlbertLayer(layer_id,layer,hidden,interventions):
     assert val.shape == hidden.shape
 
     # swap representations
-    interv_layer = interventions.pop(layer_id,None)
-    if interv_layer is not None:
-        reps = {
-                'lay': hidden,
-                'qry': qry,
-                'key': key,
-                'val': val,
-                }
-        for rep_type in ['lay','qry','key','val']:
-            interv_rep = interv_layer.pop(rep_type,None)
-            if interv_rep is not None:
-                new_state = reps[rep_type].clone()
-                for head_id, pos, swap_ids in interv_rep:
-                    new_state[swap_ids[0],pos,head_dim*head_id:head_dim*(head_id+1)] = reps[rep_type][swap_ids[1],pos,head_dim*head_id:head_dim*(head_id+1)]
-                    new_state[swap_ids[1],pos,head_dim*head_id:head_dim*(head_id+1)] = reps[rep_type][swap_ids[0],pos,head_dim*head_id:head_dim*(head_id+1)]
-                reps[rep_type] = new_state.clone()
+    reps = {
+            'lay': hidden,
+            'qry': qry,
+            'key': key,
+            'val': val,
+            }
+    for rep_type in ['lay','qry','key','val']:
+        interv_rep = interventions[layer_id][rep_type]
+        new_state = reps[rep_type].clone()
+        for head_id, pos, swap_ids in interv_rep:
+            new_state[swap_ids[0],:,head_dim*head_id:head_dim*(head_id+1)][pos,:] = reps[rep_type][swap_ids[1],:,head_dim*head_id:head_dim*(head_id+1)][pos,:]
+            new_state[swap_ids[1],:,head_dim*head_id:head_dim*(head_id+1)][pos,:] = reps[rep_type][swap_ids[0],:,head_dim*head_id:head_dim*(head_id+1)][pos,:]
+        reps[rep_type] = new_state.clone()
 
-        hidden = reps['lay'].clone()
-        qry = reps['qry'].clone()
-        key = reps['key'].clone()
-        val = reps['val'].clone()
+    hidden = reps['lay'].clone()
+    qry = reps['qry'].clone()
+    key = reps['key'].clone()
+    val = reps['val'].clone()
 
 
     #split into multiple heads