fix rope test code

Author: vpj

docs/sitemap.xml

@@ -484,7 +484,7 @@
 
     <url>
       <loc>https://nn.labml.ai/index.html</loc>
-      <lastmod>2024-01-12T16:30:00+00:00</lastmod>
+      <lastmod>2024-06-18T16:30:00+00:00</lastmod>
       <priority>1.00</priority>
     </url>
 
@@ -533,7 +533,7 @@
 
     <url>
       <loc>https://nn.labml.ai/cfr/kuhn/index.html</loc>
-      <lastmod>2023-04-02T16:30:00+00:00</lastmod>
+      <lastmod>2024-06-18T16:30:00+00:00</lastmod>
       <priority>1.00</priority>
     </url>
 
@@ -911,7 +911,7 @@
 
     <url>
       <loc>https://nn.labml.ai/transformers/models.html</loc>
-      <lastmod>2023-10-24T16:30:00+00:00</lastmod>
+      <lastmod>2024-06-18T16:30:00+00:00</lastmod>
       <priority>1.00</priority>
     </url>
 

docs/transformers/rope/index.html

@@ -436,7 +436,7 @@ <h3>Calculate scores between queries and keys</h3>
 <span class="lineno">223</span>    <span class="n">x</span> <span class="o">=</span> <span class="n">x</span><span class="p">[:,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="p">:]</span>
 <span class="lineno">224</span>    <span class="n">inspect</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
 <span class="lineno">225</span>
-<span class="lineno">226</span>    <span class="n">rotary_pe</span> <span class="o">=</span> <span class="n">RotaryPositionalEmbeddings</span><span class="p">(</span><span class="mi">3</span><span class="p">)</span>
+<span class="lineno">226</span>    <span class="n">rotary_pe</span> <span class="o">=</span> <span class="n">RotaryPositionalEmbeddings</span><span class="p">(</span><span class="mi">4</span><span class="p">)</span>
 <span class="lineno">227</span>    <span class="n">inspect</span><span class="p">(</span><span class="n">rotary_pe</span><span class="p">(</span><span class="n">x</span><span class="p">))</span>
 <span class="lineno">228</span>
 <span class="lineno">229</span>

labml_nn/transformers/rope/__init__.py

@@ -223,7 +223,7 @@ def _test_rotary():
     x = x[:, None, None, :]
     inspect(x)
 
-    rotary_pe = RotaryPositionalEmbeddings(3)
+    rotary_pe = RotaryPositionalEmbeddings(4)
     inspect(rotary_pe(x))