From 8b72ff6b74e27f3a52147aaa5cf39ac483731a02 Mon Sep 17 00:00:00 2001
From: Daniel Kukula <daniel.kuku@gmail.com>
Date: Sun, 4 Oct 2026 10:19:16 +0200
Subject: [PATCH 1/2] Fix dynamic RoPE scaling formula

---
 lib/bumblebee/layers.ex        |  6 ++---
 test/bumblebee/layers_test.exs | 45 ++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 4 deletions(-)
 create mode 100644 test/bumblebee/layers_test.exs

diff --git a/lib/bumblebee/layers.ex b/lib/bumblebee/layers.ex
index 27d81990..f4bfd087 100644
--- a/lib/bumblebee/layers.ex
+++ b/lib/bumblebee/layers.ex
@@ -1267,10 +1267,8 @@ defmodule Bumblebee.Layers do
         positions_cos_sin(position, inv_frequency)
 
       %{type: :dynamic, factor: factor} when sequence_length > max_positions ->
-        base =
-          base
-          |> Nx.multiply(factor * sequence_length / max_positions - (factor - 1))
-          |> Nx.pow(size / (size - 2))
+        scale = factor * sequence_length / max_positions - (factor - 1)
+        base = Nx.multiply(base, Nx.pow(scale, size / (size - 2)))
 
         inv_frequency = inv_frequency(base, range)
         positions_cos_sin(position, inv_frequency)
diff --git a/test/bumblebee/layers_test.exs b/test/bumblebee/layers_test.exs
new file mode 100644
index 00000000..11d84382
--- /dev/null
+++ b/test/bumblebee/layers_test.exs
@@ -0,0 +1,45 @@
+defmodule Bumblebee.LayersTest do
+  use ExUnit.Case, async: true
+
+  import Bumblebee.TestHelpers
+
+  alias Bumblebee.Layers
+
+  describe "rotary_embedding/6" do
+    test "applies dynamic scaling when sequence length exceeds max_positions" do
+      query = Axon.input("query", shape: {1, 4, 1, 4})
+      key = Axon.input("key", shape: {1, 4, 1, 4})
+      position_ids = Axon.input("position_ids", shape: {1, 4})
+      attention_mask = Axon.input("attention_mask", shape: {1, 4})
+
+      {query, _key} =
+        Layers.rotary_embedding(query, key, position_ids, attention_mask, 4,
+          max_positions: 2,
+          base: 10_000,
+          scaling_strategy: %{type: :dynamic, factor: 2.0}
+        )
+
+      model = Axon.container(%{query: query})
+
+      inputs = %{
+        "query" => Nx.broadcast(1.0, {1, 4, 1, 4}),
+        "key" => Nx.broadcast(1.0, {1, 4, 1, 4}),
+        "position_ids" => Nx.tensor([[0, 1, 2, 3]]),
+        "attention_mask" => Nx.tensor([[1, 1, 1, 1]])
+      }
+
+      {init, predict} = Axon.build(model)
+      params = init.(inputs, Axon.ModelState.empty())
+      outputs = predict.(params, inputs)
+
+      scale = 2.0 * 4 / 2 - (2.0 - 1.0)
+      base = 10_000 * :math.pow(scale, 4 / (4 - 2))
+      angle = 1 / :math.sqrt(base)
+
+      assert_all_close(
+        outputs.query[[0, 1, 0, 1]],
+        Nx.tensor(:math.cos(angle) - :math.sin(angle))
+      )
+    end
+  end
+end

From 4246c4709c734daf592b3ee4e880f92a4680ba77 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Mon, 5 Oct 2026 13:58:15 +0200
Subject: [PATCH 2/2] Use reference value from Python in dynamic RoPE test

Assisted-By: Claude Opus 5.5
---
 test/bumblebee/layers_test.exs | 9 +--------
 1 file changed, 1 insertion(+), 8 deletions(-)

diff --git a/test/bumblebee/layers_test.exs b/test/bumblebee/layers_test.exs
index 11d84382..9a07de39 100644
--- a/test/bumblebee/layers_test.exs
+++ b/test/bumblebee/layers_test.exs
@@ -32,14 +32,7 @@ defmodule Bumblebee.LayersTest do
       params = init.(inputs, Axon.ModelState.empty())
       outputs = predict.(params, inputs)
 
-      scale = 2.0 * 4 / 2 - (2.0 - 1.0)
-      base = 10_000 * :math.pow(scale, 4 / (4 - 2))
-      angle = 1 / :math.sqrt(base)
-
-      assert_all_close(
-        outputs.query[[0, 1, 0, 1]],
-        Nx.tensor(:math.cos(angle) - :math.sin(angle))
-      )
+      assert_all_close(outputs.query[[0, 1, 0, 1]], Nx.tensor(0.9967))
     end
   end
 end