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..9a07de39 --- /dev/null +++ b/test/bumblebee/layers_test.exs @@ -0,0 +1,38 @@ +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) + + assert_all_close(outputs.query[[0, 1, 0, 1]], Nx.tensor(0.9967)) + end + end +end