From 38d8744d22f64331d8bcd475bda2a11f548ab3b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luk=C3=A1=C5=A1=20Zahradn=C3=ADk?= Date: Sun, 8 Jan 2023 23:33:42 +0100 Subject: [PATCH] Fix multihead attention embedding dimensions --- neuralogic/nn/module/general/attention.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/neuralogic/nn/module/general/attention.py b/neuralogic/nn/module/general/attention.py index 12742b14..5467dbe2 100644 --- a/neuralogic/nn/module/general/attention.py +++ b/neuralogic/nn/module/general/attention.py @@ -73,9 +73,9 @@ def __init__( self.embed_dim = embed_dim self.num_heads = num_heads self.output_name = output_name - self.query_name = query_name - self.key_name = key_name - self.value_name = value_name + self.queries = query_name + self.keys = key_name + self.values = value_name self.vdim = vdim if vdim is not None else embed_dim self.kdim = kdim if kdim is not None else embed_dim self.mask_name = mask_name @@ -118,9 +118,9 @@ def __call__(self): for i in range(self.num_heads): meta = [Transformation.SLICE(rows=(i * size, (i + 1) * size))] - multihead_rules.append((q_proj(i, *terms) <= R.get(self.query_name)(terms)[q_weight:dim, dim]) | meta) - multihead_rules.append((v_proj(i, *terms) <= R.get(self.value_name)(terms)[v_weight:dim, dim]) | meta) - multihead_rules.append((k_proj(i, *terms) <= R.get(self.key_name)(terms)[k_weight:dim, dim]) | meta) + multihead_rules.append((q_proj(i, *terms) <= R.get(self.queries)(terms)[q_weight:dim, dim]) | meta) + multihead_rules.append((v_proj(i, *terms) <= R.get(self.values)(terms)[v_weight:dim, self.vdim]) | meta) + multihead_rules.append((k_proj(i, *terms) <= R.get(self.keys)(terms)[k_weight:dim, self.kdim]) | meta) attention_concat.append(R.get(attention_name)(i, *terms)) multihead_rules.append( @@ -128,11 +128,11 @@ def __call__(self): ) else: multihead_rules = [ - (q_proj(terms)[q_weight:dim, dim] <= R.get(self.query_name)(terms)) | [Transformation.IDENTITY], + (q_proj(terms)[q_weight:dim, dim] <= R.get(self.queries)(terms)) | [Transformation.IDENTITY], q_proj / self.arity | [Transformation.IDENTITY], - (v_proj(terms)[v_weight:dim, self.vdim] <= R.get(self.value_name)(terms)) | [Transformation.IDENTITY], + (v_proj(terms)[v_weight:dim, self.vdim] <= R.get(self.values)(terms)) | [Transformation.IDENTITY], v_proj / self.arity | [Transformation.IDENTITY], - (k_proj(terms)[k_weight:dim, self.kdim] <= R.get(self.key_name)(terms)) | [Transformation.IDENTITY], + (k_proj(terms)[k_weight:dim, self.kdim] <= R.get(self.keys)(terms)) | [Transformation.IDENTITY], k_proj / self.arity | [Transformation.IDENTITY], (output_rel(terms)[dim, dim] <= R.get(attention_name)(terms)) | [Transformation.IDENTITY], output_rel / self.arity | [Transformation.IDENTITY],