Check that we correctly set the key transform

danieldk · Dec 13, 2023 · 06be6dd · 06be6dd
1 parent e92d581
commit 06be6dd
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/thinc/tests/layers/test_parametric_attention_v2.py b/thinc/tests/layers/test_parametric_attention_v2.py
@@ -0,0 +1,10 @@
+from thinc.layers.parametricattention_v2 import (
+    KEY_TRANSFORM_REF,
+    ParametricAttention_v2,
+)
+from thinc.layers.gelu import Gelu
+
+
+def test_key_transform_used():
+    attn = ParametricAttention_v2(key_transform=Gelu())
+    assert attn.get_ref(KEY_TRANSFORM_REF).name == "gelu"