From dc969702ba207a06bad2144f234fb8726fc9e633 Mon Sep 17 00:00:00 2001
From: Daniel Kukula <daniel.kuku@gmail.com>
Date: Tue, 29 Sep 2026 19:00:58 +0200
Subject: [PATCH] Add bidirectional attention support to Gemma 3

---
 lib/bumblebee/text/gemma3_text.ex        |  9 ++++++--
 test/bumblebee/text/gemma3_text_test.exs | 29 ++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/lib/bumblebee/text/gemma3_text.ex b/lib/bumblebee/text/gemma3_text.ex
index 7a77997f..06bd0dfb 100644
--- a/lib/bumblebee/text/gemma3_text.ex
+++ b/lib/bumblebee/text/gemma3_text.ex
@@ -92,6 +92,10 @@ defmodule Bumblebee.Text.Gemma3Text do
       tie_word_embeddings: [
         default: true,
         doc: "whether to tie input and output embedding weights"
+      ],
+      use_bidirectional_attention: [
+        default: false,
+        doc: "whether to use bidirectional attention instead of causal attention"
       ]
     ] ++
       Shared.common_options([:rotary_embedding_scaling_strategy, :num_labels, :id_to_label]) ++
@@ -407,7 +411,7 @@ defmodule Bumblebee.Text.Gemma3Text do
             attention_head_size: spec.attention_head_size,
             attention_scale: :math.pow(spec.attention_scale_base, -0.5),
             attention_window_size: attention_window_size,
-            causal: true,
+            causal: not spec.use_bidirectional_attention,
             query_norm:
               &Layers.rms_norm(&1, shift: 1.0, epsilon: spec.layer_norm_epsilon, name: &2),
             key_norm:
@@ -535,7 +539,8 @@ defmodule Bumblebee.Text.Gemma3Text do
                  "full_attention" => :full_attention
                })
              )},
-          tie_word_embeddings: {"tie_word_embeddings", boolean()}
+          tie_word_embeddings: {"tie_word_embeddings", boolean()},
+          use_bidirectional_attention: {"use_bidirectional_attention", boolean()}
         ) ++
           Shared.rotary_embedding_options_from_transformers(data) ++
           Shared.common_options_from_transformers(data, spec)
diff --git a/test/bumblebee/text/gemma3_text_test.exs b/test/bumblebee/text/gemma3_text_test.exs
index 6af3ce6a..f1b4de39 100644
--- a/test/bumblebee/text/gemma3_text_test.exs
+++ b/test/bumblebee/text/gemma3_text_test.exs
@@ -29,6 +29,35 @@ defmodule Bumblebee.Text.Gemma3TextTest do
     )
   end
 
+  test ":base with bidirectional attention" do
+    assert {:ok, %{model: model, params: params, spec: spec}} =
+             Bumblebee.load_model({:hf, "bumblebee-testing/tiny-random-Gemma3TextModel"},
+               spec_overrides: [use_bidirectional_attention: true]
+             )
+
+    assert %Bumblebee.Text.Gemma3Text{
+             architecture: :base,
+             use_bidirectional_attention: true
+           } = spec
+
+    inputs = %{
+      "input_ids" => Nx.tensor([[10, 20, 30, 40, 50, 60, 70, 80, 0, 0]]),
+      "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
+    }
+
+    outputs = Axon.predict(model, params, inputs)
+
+    assert Nx.shape(outputs.hidden_state) == {1, 10, 32}
+
+    assert_all_close(
+      outputs.hidden_state[[.., 1..3, 1..3]],
+      Nx.tensor([
+        [[0.3755, -0.6418, 1.9320], [-0.5731, 0.3630, 2.5736], [-0.0530, 0.5178, 1.3310]]
+      ]),
+      atol: 1.0e-4
+    )
+  end
+
   test ":for_sequence_classification" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(