Query Key Value version of cross attention layer

Author: neurlang

layer/crossattention/combiner.go

@@ -8,24 +8,35 @@ func (f *CrossAttention) Put(n int, v bool) {
 // Feature returns the n-th feature from the combiner. Next layer reads
 // its inputs using this method for hashtron n in the next layer.
 func (f *CrossAttention) Feature(n int) (o uint32) {
-/*
-iov := (n % 3)
-if iov == 2 {
-return
-}
-
-dim := f.dim
-beginhead := (n / dim) * dim
+if f.qkv {
+iov := n % 3
+dim := f.dim
+beginhead := (n / dim) * dim
 
-for x := iov ^ 1; x < dim; x += 3 {
-others := f.vec[beginhead + x]
-value := f.vec[beginhead + x + (1 << iov)]
-me := f.vec[n]
-if others && me == value {
-o++
+if iov == 2 {
+// Handle the value position
+for x := 0; x < dim; x += 3 {
+query := f.vec[beginhead + x] // Get the query
+key := f.vec[beginhead + x + 1] // Get the key
+me := f.vec[n] // Current value
+if query && key && me {
+o++
+}
+}
+} else {
+// Handle query and key positions
+for x := 0; x < dim; x += 3 {
+others := f.vec[beginhead + x + iov] // Get the query or key based on iov
+value := f.vec[beginhead + x + 2] // Get the value
+me := f.vec[n] // Current query or key
+if others && me == value {
+o++
+}
+}
 }
+return o
 }
-*/
+
 iov := n & 1
 dim := f.dim
 beginhead := (n / dim) * dim

layer/crossattention/layer.go

@@ -4,16 +4,35 @@ package crossattention
 import "github.com/neurlang/classifier/layer"
 
 type CrossAttentionLayer struct {
+qkv bool
 dim int
 heads int
 }
 
 type CrossAttention struct {
 vec []bool
+qkv bool
 dim int
 heads int
 }
 
+// MustNew3 creates a new qkv full layer with size and bits
+func MustNew3(dim int, heads int) *CrossAttentionLayer {
+o, err := New3(dim, heads)
+if err != nil {
+panic(err.Error())
+}
+return o
+}
+
+// New3 creates a new qkv full layer with size and bits
+func New3(dim int, heads int) (o *CrossAttentionLayer, err error) {
+o = new(CrossAttentionLayer)
+o.dim = dim
+o.heads = heads
+o.qkv = true
+return
+}
 // MustNew creates a new full layer with size and bits
 func MustNew(dim int, heads int) *CrossAttentionLayer {
 o, err := New(dim, heads)
@@ -37,5 +56,6 @@ func (i *CrossAttentionLayer) Lay() layer.Combiner {
 o.vec = make([]bool, i.dim * i.heads)
 o.dim = i.dim
 o.heads = i.heads
+o.qkv = i.qkv
 return o
 }