update the logic for make_relative_positions

OpenNMT · Sep 11, 2024 · 4de051d · 4de051d
1 parent 2b2be50
commit 4de051d
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/layers/attention.cc b/src/layers/attention.cc
@@ -16,8 +16,8 @@ namespace ctranslate2 {
     StorageView make_relative_positions(dim_t queries_length,
                                         dim_t keys_length,
                                         dim_t max_position,
-                                        dim_t left_max_position,
-                                        dim_t right_max_position) {
+                                        dim_t left_max_position = 0,
+                                        dim_t right_max_position = 0) {
       StorageView positions({queries_length, keys_length}, DataType::INT32);
       auto* positions_data = positions.data<int32_t>();
 

diff --git a/tests/layers_test.cc b/tests/layers_test.cc
@@ -21,6 +21,16 @@ TEST(LayerTest, MakeRelativePositions2D) {
   expect_storage_eq(positions, expected);
 }
 
+TEST(LayerTest, MakeRelativePositions2D) {
+  const StorageView positions = layers::make_relative_positions(4, 4, 0, 3, 2);
+  const StorageView expected({4, 4}, std::vector<int32_t>{
+      3, 4, 5, 5,
+      2, 3, 4, 5,
+      1, 2, 3, 4,
+      0, 1, 2, 3});
+  expect_storage_eq(positions, expected);
+}
+
 TEST_P(LayerDeviceFPTest, Alibi) {
   const Device device = GetParam().device;
   const DataType dtype = GetParam().dtype;