Fix DinoV2 crash when batch_size > 1. (#429)

Co-authored-by: Jonatan Kłosko <[email protected]>

Author: zacharydenton

lib/bumblebee/vision/dino_v2.ex

@@ -312,22 +312,26 @@ defmodule Bumblebee.Vision.DinoV2 do
     Axon.layer(
       fn position_embeddings, pixel_values, _opts ->
         original_positions = div(spec.image_size, spec.patch_size)
-        {batch_size, height, width, _channels} = Nx.shape(pixel_values)
+        {_batch_size, height, width, _channels} = Nx.shape(pixel_values)
         resized_height = div(height, spec.patch_size)
         resized_width = div(width, spec.patch_size)
+        position_embeddings_batch_size = Nx.axis_size(position_embeddings, 0)
 
         class_position_embedding = position_embeddings[[.., 0..0//1, ..]]
         input_position_embeddings = position_embeddings[[.., 1..-1//1, ..]]
 
         interpolated_position_embeddings =
           input_position_embeddings
-          |> Nx.reshape({batch_size, original_positions, original_positions, spec.hidden_size})
+          |> Nx.reshape(
+            {position_embeddings_batch_size, original_positions, original_positions,
+             spec.hidden_size}
+          )
           |> Axon.Layers.resize(
             size: {resized_height, resized_width},
             method: :bicubic,
             antialias: false
           )
-          |> Nx.reshape({batch_size, :auto, spec.hidden_size})
+          |> Nx.reshape({position_embeddings_batch_size, :auto, spec.hidden_size})
 
         Nx.concatenate([class_position_embedding, interpolated_position_embeddings], axis: 1)
       end,

test/bumblebee/vision/dino_v2_test.exs

@@ -68,6 +68,59 @@ defmodule Bumblebee.Vision.DinoV2Test do
     )
   end
 
+  test ":base with batch size > 1 and position embedding interpolation" do
+    assert {:ok, %{model: model, params: params, spec: spec}} =
+             Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-Dinov2Model"})
+
+    assert %Bumblebee.Vision.DinoV2{architecture: :base} = spec
+
+    inputs = %{
+      # simulate a batch of 2 different images
+      "pixel_values" =>
+        Nx.stack([
+          Nx.broadcast(0.5, {64, 64, 3}),
+          Nx.broadcast(-0.5, {64, 64, 3})
+        ])
+    }
+
+    outputs = Axon.predict(model, params, inputs)
+
+    assert Nx.shape(outputs.hidden_state) == {2, 1025, 32}
+    assert Nx.shape(outputs.pooled_state) == {2, 32}
+
+    assert_all_close(
+      outputs.hidden_state[[0, 1..3, 1..3]],
+      Nx.tensor([
+        [-1.2287, -0.2291, -0.4323],
+        [-1.1548, -0.4430, -0.4710],
+        [-1.0547, -0.7580, -0.4654]
+      ]),
+      atol: 1.0e-1
+    )
+
+    assert_all_close(
+      outputs.pooled_state[[0, 1..3]],
+      Nx.tensor([-0.7270, -0.5913, 0.7701]),
+      atol: 1.0e-2
+    )
+
+    assert_all_close(
+      outputs.hidden_state[[1, 1..3, 1..3]],
+      Nx.tensor([
+        [0.5043, 0.9529, 0.8042],
+        [0.5555, 0.8786, 0.8492],
+        [0.6114, 0.5701, 0.9278]
+      ]),
+      atol: 1.0e-1
+    )
+
+    assert_all_close(
+      outputs.pooled_state[[1, 1..3]],
+      Nx.tensor([0.1374, -0.9757, 1.1830]),
+      atol: 1.0e-2
+    )
+  end
+
   test ":base with swiglu ffn" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(