 array, a PyTorch tensor or an attention output from Huggingface transformers.</span>
<span class="sd">        For numpy arrays or PyTorch tensors, it should have the shape `[q_len, k_len]` or `[heads, q_len, k_len]`</span>
<span class="sd">        or `[layers, heads, q_len, k_len]`.</span>
<span class="sd">    query_tokens : (List[str], optional)</span>
<span class="sd">        This is the list of query tokens.</span>
<span class="sd">    key_tokens : (List[str], optional)</span>
<span class="sd">        T