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

dl_backtrace/pytorch_backtrace/backtrace/backtrace.py

@@ -1,6 +1,7 @@
 import numpy as np
 import torch
 import torch.nn as nn
+from tqdm import tqdm
 from dl_backtrace.pytorch_backtrace.backtrace.utils import contrast as UC
 from dl_backtrace.pytorch_backtrace.backtrace.utils import prop as UP
 from dl_backtrace.pytorch_backtrace.backtrace.config import activation_master
@@ -26,6 +27,7 @@ class Backtrace(object):
             self.model_weights = EN.extract_encoder_weights(model)
             # # calculate the output of each submodule of the encoder model
             # self.all_out_model = EN.create_encoder_output(model)
+            self.activation_dict = None
             
         elif model_type == 'encoder_decoder':
             self.model = model
@@ -36,7 +38,7 @@ class Backtrace(object):
             self.model_weights = ED.extract_encoder_decoder_weights(model)  
             # # calculate the output of each submodule of the encoder-decoder model
             # self.all_out_model = ED.calculate_encoder_decoder_output(model)
-            
+            self.activation_dict = None
         
         else:
             self.model_type = model_type
@@ -317,10 +319,10 @@ class Backtrace(object):
                 all_out=all_out,
                 start_wt=start_wt,
                 multiplier=multiplier,
-                scaler=0,
+                scaler=scaler,
                 max_unit=0,
                 predicted_token=predicted_token,
-                thresholding=0.5,
+                thresholding=thresholding,
                 task="binary-classification",
             )
             return output
@@ -351,11 +353,11 @@ 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].detach().numpy())
+                start_wt = UP.calculate_start_wt(all_out[out_layer].detach().numpy(), scaler=scaler)
                 all_wt[out_layer] = start_wt * multiplier
-                layer_stack = self.layer_stack
-                all_wts = self.model_weights
-            if self.model_type == 'encoder_decoder':
+                layer_stack = self.layer_stack 
+                all_wts = self.model_weights 
+            elif self.model_type == 'encoder_decoder':
                 start_wt = UP.calculate_enc_dec_start_wt(all_out[out_layer][0].detach().numpy(), predicted_token)
                 all_wt[out_layer] = start_wt * multiplier
                 layer_stack = self.layer_stack
@@ -365,7 +367,7 @@ class Backtrace(object):
                 all_wt[out_layer] = start_wt * multiplier
                 layer_stack = self.layer_stack
                 
-        for start_layer in layer_stack:
+        for start_layer in tqdm(layer_stack):
             if model_resource[1][start_layer]["child"]:
                 child_nodes = model_resource[1][start_layer]["child"]
                 for ch in child_nodes:
@@ -538,21 +540,25 @@ class Backtrace(object):
                 elif model_resource[1][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].detach().numpy(),
                         self_attention_weights,
+                        config
                     )
                     all_wt[child_nodes[0]] += temp_wt
+                    
                 elif model_resource[1][start_layer]["class"] == 'Residual':
-                    temp_wt = UP.calculate_wt_add(
+                    temp_wt = UP.calculate_wt_residual(
                         all_wt[start_layer],
                         [all_out[ch].detach().numpy() for ch in child_nodes],
                     )
 
                     for ind, ch in enumerate(child_nodes):
                         all_wt[ch] += temp_wt[ind]
+                        
                 elif model_resource[1][start_layer]["class"] == 'Feed_Forward':
                     weights = all_wts[start_layer]
                     feed_forward_weights = HP.rename_feed_forward_keys(weights)

dl_backtrace/pytorch_backtrace/backtrace/utils/prop.py

@@ -1183,27 +1183,85 @@ def stabilize(matrix, epsilon=1e-6):
     return matrix + epsilon * np.sign(matrix)
 
 
-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_wt_residual(wts, inp=None):
+    wt_mat = []
+    inp_list = []
+    expanded_wts = as_strided(
+        wts,
+        shape=(np.prod(wts.shape),),
+        strides=(wts.strides[-1],),
+        writeable=False,  # totally use this to avoid writing to memory in weird places
+    )
+
+    for x in inp:
+        expanded_input = as_strided(
+            x,
+            shape=(np.prod(x.shape),),
+            strides=(x.strides[-1],),
+            writeable=False,  # totally use this to avoid writing to memory in weird places
+        )
+        inp_list.append(expanded_input)
+        wt_mat.append(np.zeros_like(expanded_input))
+    wt_mat = np.array(wt_mat)
+    inp_list = np.array(inp_list)
+    for i in range(wt_mat.shape[1]):
+        wt_ind1 = wt_mat[:, i]
+        wt = expanded_wts[i]
+        l1_ind1 = inp_list[:, i]
+        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
+        t_sum = p_sum - n_sum
+        p_agg_wt = 0
+        n_agg_wt = 0
+        if p_sum + n_sum > 0:
+            p_agg_wt = p_sum / (p_sum + n_sum)
+            n_agg_wt = n_sum / (p_sum + n_sum)
+        if p_sum == 0:
+            p_sum = 1
+        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[:, i] = wt_ind1
+    wt_mat = [i.reshape(wts.shape) for i in list(wt_mat)]
+    return wt_mat
+
+
+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
             n_ind = l1_ind1 < 0
             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
 
@@ -1212,21 +1270,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
 
 
-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
@@ -1234,7 +1293,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 = {
@@ -1253,12 +1326,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
 
@@ -1267,14 +1341,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
@@ -1286,28 +1406,82 @@ def calculate_wt_self_attention(wts, inp, w):
         Step-2: outputs = F.softmax(inputs, dim=dim, dtype=dtype)
         Step-3: outputs = input_a * input_b
     '''
+    # print(f"inp: {inp.shape}, wts: {wts.shape}")   # (1, 512)
+    # print(f"w['W_q']: {w['W_q'].shape}, w['W_k']: {w['W_k'].shape}, w['W_v']: {w['W_v'].shape}")
+
     query_output = np.einsum('ij,kj->ik', inp, w['W_q'])
     key_output = np.einsum('ij,kj->ik', inp, w['W_k'])
     value_output = np.einsum('ij,kj->ik', inp, w['W_v'])
 
+    # --------------- 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)
+
+    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 --------------------------------
-    wt_mat_V = calculate_relevance_V(relevance_V, value_output)
+    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
 
 
@@ -1323,7 +1497,9 @@ def calculate_wt_feed_forward(wts, inp, w):
         R2 = wts[i]
         contribution_matrix2 = np.einsum('ij,j->ij', w['W_out'], 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]
             wt_ind1 = wt_mat2[j]
@@ -1333,14 +1509,23 @@ def calculate_wt_feed_forward(wts, inp, w):
             n_ind = l1_ind1 < 0
             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
 
@@ -1359,6 +1544,9 @@ def calculate_wt_feed_forward(wts, inp, w):
         R1 = relevance_out[i]
         contribution_matrix1 = np.einsum('ij,j->ij', w['W_int'], 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]
@@ -1370,7 +1558,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 = {
@@ -1389,12 +1585,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
 
@@ -1551,7 +1748,7 @@ 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))
+    relevance_inp *= (np.sum(wts) / np.sum(relevance_inp))
     return relevance_inp
 
 

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.19'
-__version_tuple__ = version_tuple = (0, 0, 19)
+__version__ = version = '0.0.20.dev36'
+__version_tuple__ = version_tuple = (0, 0, 20, 'dev36')

{dl_backtrace-0.0.19.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.19
+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.19.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=CYabGzkNwriz1Zjt5kNvBOZD6wtqQ_twYh4s5xzmT-I,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=eeq1tj934j91bsXqZ3FbggFbBqo51bBxXWW6VuHVaMk,44153
+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=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=40fJRTsI5cvb-dbezrklSEWWghu9QobCe8TxaJQWtqU,58982
+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.19.dist-info/LICENSE,sha256=RTqAU0MFv1q3ZXKewNobKxIIPzRHgImom7e6ORV7X6o,1064
-dl_backtrace-0.0.19.dist-info/METADATA,sha256=tSJfzmAVuRLRLUjdXyJ-nBTis8L6yY-TKtV6YVBULao,7837
-dl_backtrace-0.0.19.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-dl_backtrace-0.0.19.dist-info/top_level.txt,sha256=gvGVYScJfW6c4aO5WMo4Aqa6NLEfmLK7VWXVx_GeiIk,13
-dl_backtrace-0.0.19.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.19.dist-info → dl_backtrace-0.0.20.dev36.dist-info}/WHEEL

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