right shift example comment fix

Author: vpj

docs/transformers/xl/relative_mha.html

@@ -90,7 +90,7 @@ <h1>Relative Multi-Headed Attention</h1>
             </div>
             <p> This method shifts <span ><span class="katex"><span aria-hidden="true" class="katex-html"><span class="base"><span class="strut" style="height:0.849108em;vertical-align:0em;"></span><span class="mord"><span class="mord coloredeq eqx" style=""><span class="mord mathnormal" style="">i</span></span><span class="msupsub"><span class="vlist-t"><span class="vlist-r"><span class="vlist" style="height:0.849108em;"><span style="top:-3.063em;margin-right:0.05em;"><span class="pstrut" style="height:2.7em;"></span><span class="sizing reset-size6 size3 mtight"><span class="mord mtight"><span class="mord mathnormal mtight">t</span><span class="mord mathnormal mtight">h</span></span></span></span></span></span></span></span></span></span></span></span></span> row of a matrix by <span ><span class="katex"><span aria-hidden="true" class="katex-html"><span class="base"><span class="strut" style="height:0.65952em;vertical-align:0em;"></span><span class="mord coloredeq eqx" style=""><span class="mord mathnormal" style="">i</span></span></span></span></span></span> columns.</p>
 <p>If the input is <code  class="highlight"><span></span><span class="p">[[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span> <span class="p">,</span><span class="mi">3</span><span class="p">],</span> <span class="p">[</span><span class="mi">4</span><span class="p">,</span> <span class="mi">5</span> <span class="p">,</span><span class="mi">6</span><span class="p">],</span> <span class="p">[</span><span class="mi">7</span><span class="p">,</span> <span class="mi">8</span><span class="p">,</span> <span class="mi">9</span><span class="p">]]</span></code>
-, the shifted result would be <code  class="highlight"><span></span><span class="p">[[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span> <span class="p">,</span><span class="mi">3</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="p">[</span><span class="mi">9</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">7</span><span class="p">]]</span></code>
+, the shifted result would be <code  class="highlight"><span></span><span class="p">[[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span> <span class="p">,</span><span class="mi">3</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">4</span><span class="p">,</span> <span class="mi">5</span><span class="p">],</span> <span class="p">[</span><span class="mi">6</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">7</span><span class="p">]]</span></code>
 . <em>Ideally we should mask out the lower triangle but it&#x27;s ok for our purpose</em>.</p>
 
         </div>

labml_nn/transformers/xl/relative_mha.py

@@ -25,7 +25,7 @@ def shift_right(x: torch.Tensor):
     This method shifts $i^{th}$ row of a matrix by $i$ columns.
 
     If the input is `[[1, 2 ,3], [4, 5 ,6], [7, 8, 9]]`, the shifted
-    result would be `[[1, 2 ,3], [0, 4, 5], [9, 0, 7]]`.
+    result would be `[[1, 2 ,3], [0, 4, 5], [6, 0, 7]]`.
     *Ideally we should mask out the lower triangle but it's ok for our purpose*.
     """