dl-backtrace 0.0.18__py3-none-any.whl → 0.0.20.dev36__py3-none-any.whl

dl_backtrace/tf_backtrace/backtrace/backtrace.py

@@ -192,15 +192,15 @@ class Backtrace(object):
         return temp_out
     
     def eval(self, all_out, start_wt=[], mode="default",multiplier=100.0, 
-                scaler=None, max_unit=0,thresholding=0.5,
+                scaler=0, max_unit=0,thresholding=0.5,
                 task="binary-classification",predicted_token=None):
 
         if mode=="default":
             output = self.proportional_eval(all_out=all_out,
-                                 start_wt=start_wt ,
-                                 multiplier=multiplier, 
-                                 scaler=scaler, 
-                                 max_unit=max_unit,
+                                start_wt=start_wt ,
+                                multiplier=multiplier, 
+                                scaler=scaler, 
+                                max_unit=max_unit,
                                 thresholding=thresholding,
                                 task=task,
                                 predicted_token=predicted_token)
@@ -219,7 +219,7 @@ class Backtrace(object):
             return output
 
     def proportional_eval(self, all_out, start_wt=[] ,
-                        multiplier=100.0, scaler=None, max_unit=0, 
+                        multiplier=100.0, scaler=0, max_unit=0, 
                         predicted_token=None, thresholding=0.5,
                         task="binary-classification"):
         model_resource = self.model_resource
@@ -229,7 +229,7 @@ class Backtrace(object):
         all_wt = {}
         if len(start_wt) == 0:
             if self.model_type == 'encoder':
-                start_wt = UP.calculate_start_wt(all_out[out_layer])
+                start_wt = UP.calculate_start_wt(all_out[out_layer], scaler=scaler)
                 all_wt[out_layer] = start_wt * multiplier
                 layer_stack = self.layer_stack
                 all_wts = self.model_weights
@@ -442,10 +442,12 @@ class Backtrace(object):
                 elif model_resource["graph"][start_layer]["class"] == "Self_Attention":
                     weights = all_wts[start_layer]
                     self_attention_weights = HP.rename_self_attention_keys(weights)
+                    config = self.model.config
                     temp_wt = UP.calculate_wt_self_attention(
                         all_wt[start_layer],
                         all_out[child_nodes[0]][0],
                         self_attention_weights,
+                        config
                     )
                     all_wt[child_nodes[0]] += temp_wt
                 elif model_resource["graph"][start_layer]["class"] == 'Residual':
@@ -502,10 +504,12 @@ class Backtrace(object):
                 elif model_resource["graph"][start_layer]["class"] == 'Cross_Attention':
                     weights = all_wts[start_layer]
                     cross_attention_weights = HP.rename_cross_attention_keys(weights)
+                    config = self.model.config
                     temp_wt = UP.calculate_wt_cross_attention(
                         all_wt[start_layer],
                         [all_out[ch][0] for ch in child_nodes],
                         cross_attention_weights,
+                        config
                     )
                     for ind, ch in enumerate(child_nodes):
                         all_wt[ch] += temp_wt[ind]

dl_backtrace/tf_backtrace/backtrace/utils/utils_prop.py

@@ -1161,14 +1161,17 @@ def calculate_wt_residual(wts, inp=None):
     return wt_mat
 
 
-def calculate_relevance_V(wts, value_output):
-    # Initialize wt_mat with zeros
-    wt_mat_V = np.zeros((wts.shape[0], wts.shape[1], *value_output.shape))
+def calculate_relevance_V(wts, value_output, w):
+    wt_mat_V = np.zeros(value_output.shape)
+    
+    if 'b_v' in w:
+        bias_v = w['b_v']
+    else:
+        bias_v = 0
 
     for i in range(wts.shape[0]):
         for j in range(wts.shape[1]):
             l1_ind1 = value_output
-            wt_ind1 = wt_mat_V[i, j]
             wt = wts[i, j]
 
             p_ind = l1_ind1 > 0
@@ -1176,12 +1179,21 @@ def calculate_relevance_V(wts, value_output):
             p_sum = np.sum(l1_ind1[p_ind])
             n_sum = np.sum(l1_ind1[n_ind]) * -1
 
+            if bias_v[i] > 0:
+                pbias = bias_v[i]
+                nbias = 0
+            else:
+                pbias = 0
+                nbias = bias_v[i] * -1
+
             if p_sum > 0:
-                p_agg_wt = p_sum / (p_sum + n_sum)
+                p_agg_wt = (p_sum + pbias) / (p_sum + n_sum + pbias + nbias)
+                p_agg_wt = p_agg_wt * (p_sum / (p_sum + pbias))
             else:
                 p_agg_wt = 0
             if n_sum > 0:
-                n_agg_wt = n_sum / (p_sum + n_sum)
+                n_agg_wt = (n_sum + nbias) / (p_sum + n_sum + pbias + nbias)
+                n_agg_wt = n_agg_wt * (n_sum / (n_sum + nbias))
             else:
                 n_agg_wt = 0
 
@@ -1190,21 +1202,22 @@ def calculate_relevance_V(wts, value_output):
             if n_sum == 0:
                 n_sum = 1
 
-            wt_ind1[p_ind] = (l1_ind1[p_ind] / p_sum) * wt * p_agg_wt
-            wt_ind1[n_ind] = (l1_ind1[n_ind] / n_sum) * wt * n_agg_wt * -1.0
+            wt_mat_V[p_ind] += (l1_ind1[p_ind] / p_sum) * wt * p_agg_wt
+            wt_mat_V[n_ind] += (l1_ind1[n_ind] / n_sum) * wt * n_agg_wt * -1.0
 
-    wt_mat_V = np.sum(wt_mat_V, axis=(0,1))
-    return wt_mat_V
+    return wt_mat_V 
 
 
-def calculate_relevance_QK(wts, QK_output):
-    # Initialize wt_mat with zeros
-    wt_mat_QK = np.zeros((wts.shape[0], wts.shape[1], *QK_output.shape))
+def calculate_relevance_QK(wts, QK_output, w):
+    wt_mat_QK = np.zeros(QK_output.shape)
+    
+    # Check if 'b_q' and 'b_k' exist in the weights, default to 0 if not
+    b_q = w['b_q'] if 'b_q' in w else 0
+    b_k = w['b_k'] if 'b_k' in w else 0
 
     for i in range(wts.shape[0]):
         for j in range(wts.shape[1]):
             l1_ind1 = QK_output
-            wt_ind1 = wt_mat_QK[i, j]
             wt = wts[i, j]
 
             p_ind = l1_ind1 > 0
@@ -1212,7 +1225,21 @@ def calculate_relevance_QK(wts, QK_output):
             p_sum = np.sum(l1_ind1[p_ind])
             n_sum = np.sum(l1_ind1[n_ind]) * -1
 
-            t_sum = p_sum - n_sum
+            if b_q[i] > 0 and b_k[i] > 0:
+                pbias = b_q[i] + b_k[i]
+                nbias = 0
+            elif b_q[i] > 0 and b_k[i] < 0:
+                pbias = b_q[i]
+                nbias = b_k[i] * -1
+            elif b_q[i] < 0 and b_k[i] > 0:
+                pbias = b_k[i]
+                nbias = b_q[i] * -1
+            else:
+                pbias = 0
+                nbias = b_q[i] + b_k[i]
+                nbias *= -1
+
+            t_sum = p_sum + pbias - n_sum - nbias
 
             # This layer has a softmax activation function
             act = {
@@ -1231,12 +1258,13 @@ def calculate_relevance_QK(wts, QK_output):
                         n_sum = 0
 
             if p_sum > 0:
-                p_agg_wt = p_sum / (p_sum + n_sum)
+                p_agg_wt = (p_sum + pbias) / (p_sum + n_sum + pbias + nbias)
+                p_agg_wt = p_agg_wt * (p_sum / (p_sum + pbias))
             else:
                 p_agg_wt = 0
-
             if n_sum > 0:
-                n_agg_wt = n_sum / (p_sum + n_sum)
+                n_agg_wt = (n_sum + nbias) / (p_sum + n_sum + pbias + nbias)
+                n_agg_wt = n_agg_wt * (n_sum / (n_sum + nbias))
             else:
                 n_agg_wt = 0
 
@@ -1245,14 +1273,60 @@ def calculate_relevance_QK(wts, QK_output):
             if n_sum == 0:
                 n_sum = 1
 
-            wt_ind1[p_ind] = (l1_ind1[p_ind] / p_sum) * wt * p_agg_wt
-            wt_ind1[n_ind] = (l1_ind1[n_ind] / n_sum) * wt * n_agg_wt * -1.0
+            wt_mat_QK[p_ind] += (l1_ind1[p_ind] / p_sum) * wt * p_agg_wt
+            wt_mat_QK[n_ind] += (l1_ind1[n_ind] / n_sum) * wt * n_agg_wt * -1.0
 
-    wt_mat_QK = np.sum(wt_mat_QK, axis=(0, 1))
     return  wt_mat_QK
 
 
-def calculate_wt_self_attention(wts, inp, w):
+def calculate_wt_attention_output_projection(wts, proj_output, w):
+    wt_mat_proj_output = np.zeros(proj_output.shape)
+    
+    if 'b_d' in w:
+        bias_d = w['b_d']
+    else:
+        bias_d = 0
+
+    for i in range(wts.shape[0]):
+        for j in range(wts.shape[1]):
+            l1_ind1 = proj_output
+            wt = wts[i, j]
+
+            p_ind = l1_ind1 > 0
+            n_ind = l1_ind1 < 0
+            p_sum = np.sum(l1_ind1[p_ind])
+            n_sum = np.sum(l1_ind1[n_ind]) * -1
+
+            if bias_d[i] > 0:
+                pbias = bias_d[i]
+                nbias = 0
+            else:
+                pbias = 0
+                nbias = bias_d[i] * -1
+
+            if p_sum > 0:
+                p_agg_wt = (p_sum + pbias) / (p_sum + n_sum + pbias + nbias)
+                p_agg_wt = p_agg_wt * (p_sum / (p_sum + pbias))
+            else:
+                p_agg_wt = 0
+            if n_sum > 0:
+                n_agg_wt = (n_sum + nbias) / (p_sum + n_sum + pbias + nbias)
+                n_agg_wt = n_agg_wt * (n_sum / (n_sum + nbias))
+            else:
+                n_agg_wt = 0
+
+            if p_sum == 0:
+                p_sum = 1
+            if n_sum == 0:
+                n_sum = 1
+
+            wt_mat_proj_output[p_ind] += (l1_ind1[p_ind] / p_sum) * wt * p_agg_wt
+            wt_mat_proj_output[n_ind] += (l1_ind1[n_ind] / n_sum) * wt * n_agg_wt * -1.0
+
+    return wt_mat_proj_output
+
+
+def calculate_wt_self_attention(wts, inp, w, config):
     '''
     Input:
         wts:  relevance score of the layer
@@ -1267,25 +1341,76 @@ def calculate_wt_self_attention(wts, inp, w):
     query_output = np.einsum('ij,kj->ik', inp, w['W_q'].T)
     key_output = np.einsum('ij,kj->ik', inp, w['W_k'].T)
     value_output = np.einsum('ij,kj->ik', inp, w['W_v'].T)
+    
+    # --------------- Reshape for Multi-Head Attention ----------------------
+    num_heads = getattr(config, 'num_attention_heads', getattr(config, 'num_heads', None))     # will work for BERT as well as T5/ Llama
+    hidden_size = getattr(config, 'hidden_size', getattr(config, 'd_model', None))             # will work for BERT as well as T5/Llama
+    if hasattr(config, 'num_key_value_heads'):
+        num_key_value_heads = config.num_key_value_heads
+    else:
+        num_key_value_heads = num_heads
+    head_dim = hidden_size // num_heads  # dimension of each attention head
+
+    query_states = np.einsum('thd->htd', query_output.reshape(query_output.shape[0], num_heads, head_dim))  # (num_heads, num_tokens, head_dim)
+    key_states = np.einsum('thd->htd', key_output.reshape(key_output.shape[0], num_key_value_heads, head_dim))  # (num_key_value_heads, num_tokens, head_dim)
+    value_states = np.einsum('thd->htd', value_output.reshape(value_output.shape[0], num_key_value_heads, head_dim))  # (num_key_value_heads, num_tokens, head_dim)
+    
+    # calculate how many times we need to repeat the key/value heads
+    n_rep = num_heads // num_key_value_heads
+    key_states = np.repeat(key_states, n_rep, axis=0)
+    value_states = np.repeat(value_states, n_rep, axis=0)
+
+    QK_output = np.einsum('hqd,hkd->hqk', query_states, key_states)    # (num_heads, num_tokens, num_tokens)
+    attn_weights = QK_output / np.sqrt(head_dim)
+
+    # Apply softmax along the last dimension (softmax over key dimension)
+    attn_weights = np.exp(attn_weights - np.max(attn_weights, axis=-1, keepdims=True))  # Numerically stable softmax
+    attn_weights = attn_weights / np.sum(attn_weights, axis=-1, keepdims=True)
+
+    # Weighted sum of values (num_heads, num_tokens, head_dim)
+    attn_output = np.einsum('hqk,hkl->hql', attn_weights, value_states)
+
+    transposed_attn_output = np.einsum('hqd->qhd', attn_output)
+    reshaped_attn_output = transposed_attn_output.reshape(transposed_attn_output.shape[0], num_heads * head_dim)
+
+    # Perform final linear projection (num_tokens, hidden_size)
+    final_output = np.einsum('qd,dh->qh', reshaped_attn_output, w['W_d'])
+
+    # ------------- Relevance calculation for Final Linear Projection -------------
+    wt_mat_attn_proj = calculate_wt_attention_output_projection(wts, final_output, w)
 
     # --------------- Relevance Calculation for Step-3 -----------------------
-    relevance_V = wts / 2
-    relevance_QK = wts / 2
+    # divide the relevance among `attn_weights` and `value_states`
+    wt_mat_attn_proj = wt_mat_attn_proj.reshape(-1, num_heads, head_dim)
+    wt_mat_attn_proj = np.einsum('qhd->hqd', wt_mat_attn_proj)
 
-    # --------------- Relevance Calculation for V --------------------------------
-    wt_mat_V = calculate_relevance_V(relevance_V, value_output)
+    stabilized_attn_output = stabilize(attn_output * 2)
+    norm_wt_mat_attn_proj = wt_mat_attn_proj / stabilized_attn_output
+    relevance_QK = np.einsum('htd,hbd->htb', norm_wt_mat_attn_proj, value_states) * attn_weights
+    relevance_V = np.einsum('hdt,hdb->htb', attn_weights, norm_wt_mat_attn_proj)  * value_states
 
+    # --------------- Relevance Calculation for V --------------------------------
+    relevance_V = np.einsum('hqd->qhd', relevance_V)
+    relevance_V = relevance_V.reshape(-1, num_heads * head_dim)
+    wt_mat_V = calculate_relevance_V(relevance_V, value_states, w)
+    
     # --------------- Transformed Relevance QK ----------------------------------
-    QK_output = np.einsum('ij,kj->ik', query_output, key_output)
-    wt_mat_QK = calculate_relevance_QK(relevance_QK, QK_output)
+    relevance_QK = np.einsum('hqd->qhd', relevance_QK)
+    relevance_QK = relevance_QK.reshape(-1, relevance_QK.shape[1] * relevance_QK.shape[2])
+    wt_mat_QK = calculate_relevance_QK(relevance_QK, QK_output, w)
 
     # --------------- Relevance Calculation for K and Q --------------------------------
     stabilized_QK_output = stabilize(QK_output * 2)
     norm_wt_mat_QK = wt_mat_QK / stabilized_QK_output
-    wt_mat_Q = np.einsum('ij,jk->ik', norm_wt_mat_QK, key_output) * query_output
-    wt_mat_K = np.einsum('ij,ik->kj', query_output, norm_wt_mat_QK) * key_output
+    wt_mat_Q = np.einsum('htd,hdb->htb', norm_wt_mat_QK, key_states) * query_states
+    wt_mat_K = np.einsum('htd,htb->hbd', query_states, norm_wt_mat_QK) * key_states
 
     wt_mat = wt_mat_V + wt_mat_K + wt_mat_Q
+
+    # Reshape wt_mat
+    wt_mat = np.einsum('htd->thd', wt_mat)
+    wt_mat = wt_mat.reshape(wt_mat.shape[0], wt_mat.shape[1] * wt_mat.shape[2])  # reshaped_array = array.reshape(8, 32 * 128)
+
     return wt_mat
 
 
@@ -1301,6 +1426,8 @@ def calculate_wt_feed_forward(wts, inp, w):
         R2 = wts[i]
         contribution_matrix2 = np.einsum('ij,j->ij', w['W_out'].T, intermediate_output[i])
         wt_mat2 = np.zeros(contribution_matrix2.shape)
+        
+        bias_out = w['b_out'] if 'b_out' in w else 0
 
         for j in range(contribution_matrix2.shape[0]):
             l1_ind1 = contribution_matrix2[j]
@@ -1312,13 +1439,22 @@ def calculate_wt_feed_forward(wts, inp, w):
             p_sum = np.sum(l1_ind1[p_ind])
             n_sum = np.sum(l1_ind1[n_ind]) * -1
 
+            # Handle positive and negative bias contributions
+            if bias_out[i] > 0:
+                pbias = bias_out[i]
+                nbias = 0
+            else:
+                pbias = 0
+                nbias = -bias_out[i]
+
             if p_sum > 0:
-                p_agg_wt = p_sum / (p_sum + n_sum)
+                p_agg_wt = (p_sum + pbias) / (p_sum + n_sum + pbias + nbias)
+                p_agg_wt = p_agg_wt * (p_sum / (p_sum + pbias))
             else:
                 p_agg_wt = 0
-
             if n_sum > 0:
-                n_agg_wt = n_sum / (p_sum + n_sum)
+                n_agg_wt = (n_sum + nbias) / (p_sum + n_sum + pbias + nbias)
+                n_agg_wt = n_agg_wt * (n_sum / (n_sum + nbias))
             else:
                 n_agg_wt = 0
 
@@ -1337,6 +1473,9 @@ def calculate_wt_feed_forward(wts, inp, w):
         R1 = relevance_out[i]
         contribution_matrix1 = np.einsum('ij,j->ij', w['W_int'].T, inp[i])
         wt_mat1 = np.zeros(contribution_matrix1.shape)
+        
+        # Check if bias 'b_int' exists, default to 0 if not
+        bias_int = w['b_int'] if 'b_int' in w else 0
 
         for j in range(contribution_matrix1.shape[0]):
             l1_ind1 = contribution_matrix1[j]
@@ -1348,7 +1487,15 @@ def calculate_wt_feed_forward(wts, inp, w):
             p_sum = np.sum(l1_ind1[p_ind])
             n_sum = np.sum(l1_ind1[n_ind]) * -1
 
-            t_sum = p_sum - n_sum
+            # Handle positive and negative bias
+            if bias_int[i] > 0:
+                pbias = bias_int[i]
+                nbias = 0
+            else:
+                pbias = 0
+                nbias = -bias_int[i]
+
+            t_sum = p_sum + pbias - n_sum - nbias
 
             # This layer has a ReLU activation function
             act = {
@@ -1367,12 +1514,13 @@ def calculate_wt_feed_forward(wts, inp, w):
                         n_sum = 0
 
             if p_sum > 0:
-                p_agg_wt = p_sum / (p_sum + n_sum)
+                p_agg_wt = (p_sum + pbias) / (p_sum + n_sum + pbias + nbias)
+                p_agg_wt = p_agg_wt * (p_sum / (p_sum + pbias))
             else:
                 p_agg_wt = 0
-
             if n_sum > 0:
-                n_agg_wt = n_sum / (p_sum + n_sum)
+                n_agg_wt = (n_sum + nbias) / (p_sum + n_sum + pbias + nbias)
+                n_agg_wt = n_agg_wt * (n_sum / (n_sum + nbias))
             else:
                 n_agg_wt = 0
 
@@ -1461,8 +1609,8 @@ def calculate_wt_pooler(wts, inp, w):
         # Calculate relevance for each token
         relevance_inp[i] = wt_mat.sum(axis=0)
 
-    relevance_inp *= (100 / np.sum(relevance_inp))
-    return relevance_inp
+    relevance_inp *= (np.sum(wts) / np.sum(relevance_inp))
+    return relevance_inp 
 
 
 def calculate_wt_classifier(wts, inp, w):
@@ -1595,7 +1743,7 @@ def calculate_wt_lm_head(wts, inp, w):
     return relevance_input
 
 
-def calculate_wt_cross_attention(wts, inp, w):
+def calculate_wt_cross_attention(wts, inp, w, config):
     '''
     Input:
         wts:  relevance score of the layer
@@ -1613,23 +1761,77 @@ def calculate_wt_cross_attention(wts, inp, w):
     key_output = np.einsum('ij,kj->ik', k_v_inp, w['W_k'].T)
     value_output = np.einsum('ij,kj->ik', k_v_inp, w['W_v'].T)
 
+    # --------------- Reshape for Multi-Head Attention ----------------------
+    num_heads = getattr(config, 'num_attention_heads', getattr(config, 'num_heads', None))     # will work for BERT as well as T5/ Llama
+    hidden_size = getattr(config, 'hidden_size', getattr(config, 'd_model', None))             # will work for BERT as well as T5/Llama
+    if hasattr(config, 'num_key_value_heads'):
+        num_key_value_heads = config.num_key_value_heads
+    else:
+        num_key_value_heads = num_heads
+    head_dim = hidden_size // num_heads  # dimension of each attention head
+
+    query_states = np.einsum('thd->htd', query_output.reshape(query_output.shape[0], num_heads, head_dim))  # (num_heads, num_tokens, head_dim)
+    key_states = np.einsum('thd->htd', key_output.reshape(key_output.shape[0], num_key_value_heads, head_dim))  # (num_key_value_heads, num_tokens, head_dim)
+    value_states = np.einsum('thd->htd', value_output.reshape(value_output.shape[0], num_key_value_heads, head_dim))  # (num_key_value_heads, num_tokens, head_dim)
+    
+    # calculate how many times we need to repeat the key/value heads
+    n_rep = num_heads // num_key_value_heads
+    key_states = np.repeat(key_states, n_rep, axis=0)
+    value_states = np.repeat(value_states, n_rep, axis=0)
+
+    QK_output = np.einsum('hqd,hkd->hqk', query_states, key_states)    # (num_heads, num_tokens, num_tokens)
+    attn_weights = QK_output / np.sqrt(head_dim)
+
+    # Apply softmax along the last dimension (softmax over key dimension)
+    attn_weights = np.exp(attn_weights - np.max(attn_weights, axis=-1, keepdims=True))  # Numerically stable softmax
+    attn_weights = attn_weights / np.sum(attn_weights, axis=-1, keepdims=True)
+
+    # Weighted sum of values (num_heads, num_tokens, head_dim)
+    attn_output = np.einsum('hqk,hkl->hql', attn_weights, value_states)
+
+    transposed_attn_output = np.einsum('hqd->qhd', attn_output)
+    reshaped_attn_output = transposed_attn_output.reshape(transposed_attn_output.shape[0], num_heads * head_dim)
+
+    # Perform final linear projection (num_tokens, hidden_size)
+    final_output = np.einsum('qd,dh->qh', reshaped_attn_output, w['W_d'])
+
+    # ------------- Relevance calculation for Final Linear Projection -------------
+    wt_mat_attn_proj = calculate_wt_attention_output_projection(wts, final_output)
+
     # --------------- Relevance Calculation for Step-3 -----------------------
-    relevance_V = wts / 2
-    relevance_QK = wts / 2
+    # divide the relevance among `attn_weights` and `value_states`
+    wt_mat_attn_proj = wt_mat_attn_proj.reshape(-1, num_heads, head_dim)
+    wt_mat_attn_proj = np.einsum('qhd->hqd', wt_mat_attn_proj)
 
-    # --------------- Relevance Calculation for V --------------------------------
-    wt_mat_V = calculate_relevance_V(relevance_V, value_output)
+    stabilized_attn_output = stabilize(attn_output * 2)
+    norm_wt_mat_attn_proj = wt_mat_attn_proj / stabilized_attn_output
+    relevance_QK = np.einsum('htd,hbd->htb', norm_wt_mat_attn_proj, value_states) * attn_weights
+    relevance_V = np.einsum('hdt,hdb->htb', attn_weights, norm_wt_mat_attn_proj)  * value_states
 
+    # --------------- Relevance Calculation for V --------------------------------
+    relevance_V = np.einsum('hqd->qhd', relevance_V)
+    relevance_V = relevance_V.reshape(-1, num_heads * head_dim)
+    wt_mat_V = calculate_relevance_V(relevance_V, value_states)
+    
     # --------------- Transformed Relevance QK ----------------------------------
-    QK_output = np.einsum('ij,kj->ik', query_output, key_output)
+    relevance_QK = np.einsum('hqd->qhd', relevance_QK)
+    relevance_QK = relevance_QK.reshape(-1, relevance_QK.shape[1] * relevance_QK.shape[2])
     wt_mat_QK = calculate_relevance_QK(relevance_QK, QK_output)
 
     # --------------- Relevance Calculation for K and Q --------------------------------
     stabilized_QK_output = stabilize(QK_output * 2)
     norm_wt_mat_QK = wt_mat_QK / stabilized_QK_output
-    wt_mat_Q = np.einsum('ij,jk->ik', norm_wt_mat_QK, key_output) * query_output
-    wt_mat_K = np.einsum('ij,ik->kj', query_output, norm_wt_mat_QK) * key_output
+    wt_mat_Q = np.einsum('htd,hdb->htb', norm_wt_mat_QK, key_states) * query_states
+    wt_mat_K = np.einsum('htd,htb->hbd', query_states, norm_wt_mat_QK) * key_states
 
+    # Relevance of KV input
     wt_mat_KV = wt_mat_V + wt_mat_K
+
+    # Reshape wt_mat_Q and wt_mat_KV
+    wt_mat_Q = np.einsum('htd->thd', wt_mat_Q)
+    wt_mat_KV = np.einsum('htd->thd', wt_mat_KV)
+    wt_mat_Q = wt_mat_Q.reshape(wt_mat_Q.shape[0], wt_mat_Q.shape[1] * wt_mat_Q.shape[2])
+    wt_mat_KV = wt_mat_KV.reshape(wt_mat_KV.shape[0], wt_mat_KV.shape[1] * wt_mat_KV.shape[2])
+
     wt_mat = [wt_mat_KV, wt_mat_Q]
     return wt_mat

dl_backtrace/version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.0.18'
-__version_tuple__ = version_tuple = (0, 0, 18)
+__version__ = version = '0.0.20.dev36'
+__version_tuple__ = version_tuple = (0, 0, 20, 'dev36')

{dl_backtrace-0.0.18.dist-info → dl_backtrace-0.0.20.dev36.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: dl_backtrace
-Version: 0.0.18
+Version: 0.0.20.dev36
 Summary: A python SDK for Deep Learning Backtrace
 Home-page: https://xai.arya.ai/docs/introduction
 License: MIT

{dl_backtrace-0.0.18.dist-info → dl_backtrace-0.0.20.dev36.dist-info}/RECORD

@@ -1,5 +1,5 @@
 dl_backtrace/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dl_backtrace/version.py,sha256=k900Q8XjzRKO6ZOHY0wFLzfzTGArI0sGircauDDJhu0,413
+dl_backtrace/version.py,sha256=tYHVV4mIeOCumN5OCSD_xV6vt2LqOBp8qrLBhN4xnyw,428
 dl_backtrace/old_backtrace/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 dl_backtrace/old_backtrace/pytorch_backtrace/__init__.py,sha256=TDhKQIj1INyQq7cqTvpjpnBDhoMeWSoqVdx_mPAV3Sw,24
 dl_backtrace/old_backtrace/pytorch_backtrace/backtrace/__init__.py,sha256=AuR7uMTbf7rrFl-9sMIQ3lQmhu_ATSFBZmfs7R66jAc,76
@@ -17,18 +17,18 @@ dl_backtrace/old_backtrace/tf_backtrace/backtrace/utils/contrast.py,sha256=p-0zj
 dl_backtrace/old_backtrace/tf_backtrace/backtrace/utils/prop.py,sha256=-h7nHEvsoEwfksHsT52VfnZy334DvqqP8g6fMFVNnAM,25670
 dl_backtrace/pytorch_backtrace/__init__.py,sha256=TDhKQIj1INyQq7cqTvpjpnBDhoMeWSoqVdx_mPAV3Sw,24
 dl_backtrace/pytorch_backtrace/backtrace/__init__.py,sha256=AuR7uMTbf7rrFl-9sMIQ3lQmhu_ATSFBZmfs7R66jAc,76
-dl_backtrace/pytorch_backtrace/backtrace/backtrace.py,sha256=wU5J7QkTnkQ0Jri8Xe6WctqeKFO4hTcL8MQpLrXbvdY,35211
+dl_backtrace/pytorch_backtrace/backtrace/backtrace.py,sha256=bdb8eFzTE7ms34ne-VBJxXY8AMHhu-vTenUNwCMKbRo,44410
 dl_backtrace/pytorch_backtrace/backtrace/config.py,sha256=ODrgOC74ojzGLnEto_ah-WPA8_MyRE7cZkZlU15239s,983
 dl_backtrace/pytorch_backtrace/backtrace/utils/__init__.py,sha256=KffAJVu7NsgfMHEZaY7lND2LQwZamVIquqx8POwOaLg,120
-dl_backtrace/pytorch_backtrace/backtrace/utils/contrast.py,sha256=owsncnnz-j7UameRxx5uL9Q1AqOcA-uxhpyyfs0DqBw,32228
+dl_backtrace/pytorch_backtrace/backtrace/utils/contrast.py,sha256=gcvDRZstcgAkjzyyxy6dagSV2EDFGqh9_SMn9P1-2lo,52795
 dl_backtrace/pytorch_backtrace/backtrace/utils/encoder.py,sha256=1dUKEL_LAuFcUEldYdyQVsV_P-KH3gTWEOgROfXPwyc,7469
 dl_backtrace/pytorch_backtrace/backtrace/utils/encoder_decoder.py,sha256=fxbXwUl7KQKwWSdKMO-bYvDC7NDct71F_U7uZzh1orw,25658
 dl_backtrace/pytorch_backtrace/backtrace/utils/helper.py,sha256=IG0XkPEfbPpWBm4aVHgv-GkgFFrl2wu5xMrz6DeH_xQ,3512
-dl_backtrace/pytorch_backtrace/backtrace/utils/prop.py,sha256=hHYFLGR3hsGtxlzdA4qXhgFFsBqq5AhCvxRUxx6KwZY,41284
+dl_backtrace/pytorch_backtrace/backtrace/utils/prop.py,sha256=oEwWmRsLwoYdCXJaEEW7-UL1CLVOsseUAcOasTkyQeE,67294
 dl_backtrace/tf_backtrace/__init__.py,sha256=TDhKQIj1INyQq7cqTvpjpnBDhoMeWSoqVdx_mPAV3Sw,24
 dl_backtrace/tf_backtrace/backtrace/__init__.py,sha256=KkU7X_wXxwYR4HYQqQ3kWtxlJK3Ytaa84e1Jbzc2_ZA,84
 dl_backtrace/tf_backtrace/backtrace/activation_info.py,sha256=3Ppw4_6rJV16YPXnKjd2WPaULgUUL6U01bh8Foa-3Yg,1334
-dl_backtrace/tf_backtrace/backtrace/backtrace.py,sha256=eQa7wz3MsfjkgZE72voKGlZGjxbf42kHkXcwB7Ve3qI,41474
+dl_backtrace/tf_backtrace/backtrace/backtrace.py,sha256=k6irwiOu68EpRG_W9wk5UF7QS50TKM5JqUQfj4qmvWw,41635
 dl_backtrace/tf_backtrace/backtrace/models.py,sha256=wPaeRuEvZL2xTdj6I6hF-ZCKC2c8EAuF9PoOIZkxkR4,466
 dl_backtrace/tf_backtrace/backtrace/server.py,sha256=jphibvI46QpQcqnpXVIYFq_M2CtRVONgItZe9iWOd54,567
 dl_backtrace/tf_backtrace/backtrace/utils/__init__.py,sha256=ci_RAYYnqyAWa_rcIEycnqCghQ4aZtvaGQ7oDUb_k_0,131
@@ -36,9 +36,9 @@ dl_backtrace/tf_backtrace/backtrace/utils/encoder.py,sha256=WeGLjIRHNqjwIK-8UB0x
 dl_backtrace/tf_backtrace/backtrace/utils/encoder_decoder.py,sha256=qbS34WswNiT1xgON5ayNAIewJMRfDdeGel6l1XrjXms,24247
 dl_backtrace/tf_backtrace/backtrace/utils/helper.py,sha256=QB21kPB5iJfRpy8khYJnzojaKf5ACnAFYh5XxYBcnXA,3419
 dl_backtrace/tf_backtrace/backtrace/utils/utils_contrast.py,sha256=rQwManW0d6Td6V_A1qGezcZ19Tgr34zbFcieQ0rqAAc,48415
-dl_backtrace/tf_backtrace/backtrace/utils/utils_prop.py,sha256=dC6QxbwpA5JhE8009vnZVQyY8eIE3RFx6t7I-N17-k0,58320
-dl_backtrace-0.0.18.dist-info/LICENSE,sha256=RTqAU0MFv1q3ZXKewNobKxIIPzRHgImom7e6ORV7X6o,1064
-dl_backtrace-0.0.18.dist-info/METADATA,sha256=YuIoncn6l2OscEaxrC5VA7OaFre9ddTh3pFGqZPJzsI,7837
-dl_backtrace-0.0.18.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
-dl_backtrace-0.0.18.dist-info/top_level.txt,sha256=gvGVYScJfW6c4aO5WMo4Aqa6NLEfmLK7VWXVx_GeiIk,13
-dl_backtrace-0.0.18.dist-info/RECORD,,
+dl_backtrace/tf_backtrace/backtrace/utils/utils_prop.py,sha256=kduJpbN2FIrrhgrUxSFkX3gnTfXLuQa4qx-LH0dDB7A,68291
+dl_backtrace-0.0.20.dev36.dist-info/LICENSE,sha256=RTqAU0MFv1q3ZXKewNobKxIIPzRHgImom7e6ORV7X6o,1064
+dl_backtrace-0.0.20.dev36.dist-info/METADATA,sha256=H5E00w4t5F6O67roa2Y6tiGEW_SziKjsgMssLEoDM_U,7843
+dl_backtrace-0.0.20.dev36.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+dl_backtrace-0.0.20.dev36.dist-info/top_level.txt,sha256=gvGVYScJfW6c4aO5WMo4Aqa6NLEfmLK7VWXVx_GeiIk,13
+dl_backtrace-0.0.20.dev36.dist-info/RECORD,,

{dl_backtrace-0.0.18.dist-info → dl_backtrace-0.0.20.dev36.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (73.0.1)
+Generator: setuptools (75.5.0)
 Root-Is-Purelib: true
 Tag: py3-none-any