@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

package/{src/layers/multihead_attention.ts → dist/layers/multihead_attention.js}

@@ -1,86 +1,58 @@
 import * as tf from '@tensorflow/tfjs';
-import { type LayerArgs } from '@tensorflow/tfjs-layers/dist/engine/topology';
-import { type Kwargs } from '@tensorflow/tfjs-layers/dist/types';
-import { generateCausalAttentionMask } from "@/utils";
-
-
-export interface MultiHeadAttentionArgs extends LayerArgs {
-    numHeads: number;
-    embedDim: number;
-    useBias?: boolean;
-    dropout?: number;
-    causal?: boolean;
-}
-
-
-export interface ScaledDotProductionAttentionKwargs {
-    training?: boolean;
-    dropout?: number;
-    causal?: boolean;
-    scaling_factor?: number;
-}
-
-
+import { causal as generateCausalMask } from "../masks";
 /**
  * This MultiHead Attention layer implements the algorithm as described in
  * the paper "Attention is all you Need" Vaswani et al., 2017.
- * 
+ *
  * @param numHeads number of attention heads to use
  * @param embedDim the embedding size of the input (input embeddings, typically the last dimension)
  * @param causal use causal masking, default `false`
  * @param dropout use dropout during the attention calculations, default `0.0`
  * @param useBias use bias for the dense sublayers, default `true`
- * 
+ *
  * The TensorFlow version uses tf.einsum, whose gradient op has not yet been
  * implemented (https://github.com/tensorflow/tfjs/pull/4955#discussion_r619219334),
  * therefore we follow the PyTorch implementation described in:
  * https://docs.pytorch.org/tutorials/intermediate/transformer_building_blocks.html#multiheadattention
  * https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
- * 
+ *
  * This implementation is different from TensorFlow's whose attention weights
  * are shaped [embed dim, heads, embed dim] where as PyTorch and OpenAI's attention weights
  * are shaped [embed dim, embed dim]
  * https://github.com/pytorch/pytorch/blob/134179474539648ba7dee1317959529fbd0e7f89/torch/nn/modules/activation.py#L1080
  * https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/model.py#L53
- * 
+ *
  * TODO: implement a fast track for self attention (query = key = value)
  * where a single dense layer combines and replaces the query, key and projection layers
- * 
+ *
  * TODO: add kDim and vDim to accept key and values whose embedding dimensions differ from query's.
  */
 export class MultiHeadAttention extends tf.layers.Layer {
     static className = "MultiHeadAttention";
-    protected readonly numHeads: number;
-    protected readonly embedDim: number; // size of embedding dim of inputs, also per attention head
-    protected readonly useBias: boolean;
-    protected readonly dropout: number;
-    protected readonly causal: boolean; // use causal attention to mask future words
-
+    numHeads;
+    embedDim; // size of embedding dim of inputs, also per attention head
+    useBias;
+    dropout;
+    causal; // use causal attention to mask future words
     // projection simply means matrix multiplying query, key, and value
     // with weights to create a representation of the inputs
-    protected readonly queryProjection: tf.layers.Layer;
-    protected readonly keyProjection: tf.layers.Layer;
-    protected readonly valueProjection: tf.layers.Layer;
-    protected readonly outputProjection: tf.layers.Layer;
-
-
-    constructor({ numHeads, embedDim, useBias = true, dropout = 0.0, causal = false, ...args }: MultiHeadAttentionArgs) {
+    queryProjection;
+    keyProjection;
+    valueProjection;
+    outputProjection;
+    constructor({ numHeads, embedDim, useBias = true, dropout = 0.0, causal = false, ...args }) {
         super(args);
-
         if (embedDim % numHeads != 0) {
             throw Error(`${this.getClassName()}::constructor ${this.name} embedDim (${embedDim}) is not divisible by numHeads (${numHeads})`);
         }
-
         this.numHeads = numHeads;
         this.embedDim = embedDim;
         this.useBias = useBias;
         this.dropout = dropout;
         this.causal = causal;
-
         if (this.dropout >= 1) {
             throw Error(`${this.getClassName()}::constructor dropout must be within [0, 1)`);
         }
-
         // intialize the projection weights, this should be in the
         // build() function but is done here to avoid linting complaints
         this.queryProjection = tf.layers.dense({ useBias, units: embedDim });
@@ -88,188 +60,134 @@ export class MultiHeadAttention extends tf.layers.Layer {
         this.valueProjection = tf.layers.dense({ useBias, units: embedDim });
         this.outputProjection = tf.layers.dense({ useBias, units: embedDim });
     }
-
-
     /**
      * Forward propagation. Provide one input tensor or three identical tensors to self-attention.
      * @param inputs a single tensor for self-attention or an array of exactly three
      *   tensors that are either identical (self-attention) or different (cross-attention)
      * @param kwargs.packingMask a mask to prevent tokens from attending across document boundaries
      */
-    override call(
-        inputs: tf.Tensor | tf.Tensor[],
-        kwargs: Kwargs & {
-            packingMask?: tf.Tensor,
-            causalMask?: tf.Tensor,
-        }
-    ): tf.Tensor | tf.Tensor[] {
+    call(inputs, kwargs) {
         // validate the input tensors
         if (!Array.isArray(inputs)) {
             inputs = [inputs];
         }
-
         // accept only 1 input (self attention) or 3 inputs (self or cross attention)
         if (inputs.length != 1 && inputs.length != 3) {
             throw Error(`${this.getClassName()}::call ${this.name} expects exactly one or three input tensors, ${inputs.length} were provided`);
         }
-
         for (const input of inputs) {
             if (input.shape.length != 3) {
                 throw Error(`${this.getClassName()}::call ${this.name} expected input shapes of [batch, seq, embed_dim], got ${JSON.stringify(input.shape)}`);
             }
         }
-
         const [query, key, value] = inputs;
         const packingMask = kwargs.packingMask ?? null;
         const causalMask = kwargs.causalMask ?? null;
-
         return inputs.length == 3
             // cross-attention
-            ? this.forward(query!, key!, value!, packingMask, causalMask, kwargs)
+            ? this.forward(query, key, value, packingMask, causalMask, kwargs)
             // self-attention
-            : this.forward(query!, query!, query!, packingMask, causalMask, kwargs);
+            : this.forward(query, query, query, packingMask, causalMask, kwargs);
     }
-
-
     /**
      * Forward propagation
      */
-    protected forward(
-        query_input: tf.Tensor,
-        key_input: tf.Tensor,
-        value_input: tf.Tensor,
-        packing_mask: tf.Tensor | null,
-        causal_mask: tf.Tensor | null,
-        kwargs: Kwargs): tf.Tensor {
-
+    forward(query_input, key_input, value_input, packing_mask, causal_mask, kwargs) {
         // dimensions abbreviations
         // batch = the number of sequences in the input
         // seq = the length of each sequence in the input
         // dims = the size of each token's embedding
         return tf.tidy(() => {
             const { query, key, value } = this.applyInputProjections(query_input, key_input, value_input);
-
             // swap the seq and heads dimensions: [batch, seq, heads, head_dim] -> [batch, heads, seq, head_dim]
             const move_head_dim_forward = [0, 2, 1, 3];
-
-            const {
-                query_split, key_split, value_split
-            } = this.splitHeads(query, key, value, move_head_dim_forward);
-
+            const { query_split, key_split, value_split } = this.splitHeads(query, key, value, move_head_dim_forward);
             // apply scaled dot production attention to get [batch, seq, numHeads, embedDim]
-            const spda = MultiHeadAttention.scaledDotProductionAttention(
-                query_split, key_split, value_split,
-                kwargs.attentionMask ?? null, packing_mask, causal_mask,
-                this.dropout, this.causal, kwargs);
-
+            const spda = MultiHeadAttention.scaledDotProductionAttention(query_split, key_split, value_split, kwargs.attentionMask ?? null, packing_mask, causal_mask, this.dropout, this.causal, kwargs);
             // concat heads and apply the output projection
-            const output = this.outputProjection.apply(
-                spda.transpose(move_head_dim_forward).reshape([query_input.shape[0], -1, this.embedDim]));
-
-            return output as tf.Tensor;
-        })
+            const output = this.outputProjection.apply(spda.transpose(move_head_dim_forward).reshape([query_input.shape[0], -1, this.embedDim]));
+            return output;
+        });
     }
-
-
-    protected applyInputProjections(query_input: tf.Tensor, key_input: tf.Tensor, value_input: tf.Tensor) {
+    applyInputProjections(query_input, key_input, value_input) {
         // apply input projections, this is a batched matrix multiplication operated on the last
         // dimension of query_input and first dimension of the dense layer weights,
         // [batch, seq, dims] x [dims, dims] = [batch x seq, dims] x [dims, dims] = [batch x seq, dims] = [batch, seq, dims]
         return tf.tidy(() => {
             return {
-                query: this.queryProjection.apply(query_input) as tf.Tensor,
-                key: this.keyProjection.apply(key_input) as tf.Tensor,
-                value: this.valueProjection.apply(value_input) as tf.Tensor
-            }
-        })
+                query: this.queryProjection.apply(query_input),
+                key: this.keyProjection.apply(key_input),
+                value: this.valueProjection.apply(value_input)
+            };
+        });
     }
-
-
-    protected splitHeads(query: tf.Tensor, key: tf.Tensor, value: tf.Tensor, shuffle: number[]) {
+    splitHeads(query, key, value, shuffle) {
         // split heads and prepare for scaled dot product attention by splitting the
         // last dimension to get the heads, bring the heads forward
         // [batch, seq, dims] -> [batch, seq, heads, dims / heads] -> [batch, heads, seq, head_dim]
         const batch_size = query.shape[0];
         const split_heads = [batch_size, -1, this.numHeads, this.embedDim / this.numHeads];
-
         return tf.tidy(() => {
             return {
-                query_split: query.reshape(split_heads).transpose(shuffle) as tf.Tensor4D,
-                key_split: key.reshape(split_heads).transpose(shuffle) as tf.Tensor4D,
-                value_split: value.reshape(split_heads).transpose(shuffle) as tf.Tensor4D
-            }
-        })
+                query_split: query.reshape(split_heads).transpose(shuffle),
+                key_split: key.reshape(split_heads).transpose(shuffle),
+                value_split: value.reshape(split_heads).transpose(shuffle)
+            };
+        });
     }
-
-
     /**
      * Applies the scaled dot-product formula: softmax(QK_t / sqrt(d_k))V,
      * formula (1) of the 2017 paper Attention Is All You Need
-     * 
-     * @param attentionMask a mask to prevent tokens from being 
+     *
+     * @param attentionMask a mask to prevent tokens from being
      *   attended to (usually for padding tokens). It should have the shape
      *   [batch, head, query_sequence_len, key_sequence_len]. To use in
      *   conjunction with causal masking, the tensor should be a boolean type
      *   where false indicates a masked token.
      * @param packingMask a mask to prevent tokens from attending across document boundaries
      */
-    static scaledDotProductionAttention(
-        query: tf.Tensor,
-        key: tf.Tensor,
-        value: tf.Tensor,
-        attentionMask: tf.Tensor | null,
-        packingMask: tf.Tensor | null,
-        causalMask: tf.Tensor | null,
-        dropout: number,
-        causal: boolean,
-        kwargs: ScaledDotProductionAttentionKwargs = {}
-    ): tf.Tensor {
+    static scaledDotProductionAttention(query, key, value, attentionMask, packingMask, causalMask, dropout, causal, kwargs = {}) {
         return tf.tidy(() => {
             const { training = false, scaling_factor } = kwargs;
-
             key.shape.forEach((val, index) => {
                 if (key.shape[index] != value.shape[index]) {
                     throw Error(`scaledDotProductionAttention: expected key and value` +
                         ` to have the same shape, got ${JSON.stringify(key.shape)} (key) and` +
                         ` ${JSON.stringify(value.shape)} (value)`);
                 }
-            })
-
-
+            });
             // mask's shape is [..., seq, seq] where seq is the number of words/tokens in the input,
             // not adding the batch dimension yet to lessen the calculations
             const causal_mask_shape = [
                 query.shape[query.shape.length - 2],
-                key.shape[key.shape.length - 2]];
-
+                key.shape[key.shape.length - 2]
+            ];
             let mask = tf.zeros(causal_mask_shape);
-
             if (causal && causal_mask_shape[0] > 1) {
                 if (attentionMask && attentionMask.dtype != "bool") {
                     throw Error(`scaledDotProductionAttention: the attention mask must be undefined or a boolean type if used with causal attention`);
                 }
-
                 // apply a causal attention mask so that tokens can only attend to preceding tokens,
                 // prevents looking at head
                 if (causalMask) {
                     mask = causalMask;
-                } else {
-                    mask = generateCausalAttentionMask(causal_mask_shape[0], causal_mask_shape[1]);
+                }
+                else {
+                    mask = generateCausalMask(causal_mask_shape[0], causal_mask_shape[1]);
                 }
             }
-
             if (attentionMask) {
                 if (attentionMask.dtype == "bool") {
                     // convert the boolean mask to float
                     // warning: do not use 1e9, it will overflow, use something smaller like 1e7
                     mask = mask.add(attentionMask.cast("float32").sub(1).mul(1e7));
-                } else {
+                }
+                else {
                     // this will occur only when not using causal masking,
                     // if the attention mask is not boolean, it's assumed the masking is already calculated,
                     mask = attentionMask;
                 }
             }
-
             // 1. matrix multiply query and transposed key
             // 2. divide by scaling factor
             // 3. apply softmax to the result
@@ -280,42 +198,33 @@ export class MultiHeadAttention extends tf.layers.Layer {
                 .matMul(key, false, true)
                 .div(Math.sqrt(scaling_factor ?? key.shape[key.shape.length - 1]))
                 .add(mask);
-
             if (packingMask) {
                 // packing mask is added separately because each mask within a batch may be different,
                 // so it cannot be broadcasted
                 pre_softmax = pre_softmax.add(packingMask);
             }
-
             const spda = tf.softmax(pre_softmax);
-
             const spda_dropout = tf.dropout(spda, training ? dropout : 0);
             const attention = spda_dropout.matMul(value);
-
             return attention;
         });
     }
-
-
-    override build(inputShape: tf.Shape | tf.Shape[]): void {
-        let input_shape: tf.Shape[] = [];
-
+    build(inputShape) {
+        let input_shape = [];
         if (Array.isArray(inputShape) && Array.isArray(inputShape[0])) {
-            input_shape = inputShape as tf.Shape[];
-        } else {
-            input_shape = [inputShape as tf.Shape, inputShape as tf.Shape, inputShape as tf.Shape];
+            input_shape = inputShape;
+        }
+        else {
+            input_shape = [inputShape, inputShape, inputShape];
         }
-
         if (input_shape.length != 1 && input_shape.length != 3) {
             throw Error(`${this.getClassName()}::build ${this.name} accepts either exactly one or three inputs, received ${JSON.stringify(inputShape)}`);
         }
-
         // initialize the sublayer weights
         this.queryProjection.build(input_shape[0]);
         this.keyProjection.build(input_shape[1]);
         this.valueProjection.build(input_shape[2]);
         this.outputProjection.build(input_shape[0]);
-
         // the sublayer weights need to be tracked by this layer otherwise
         // backpropagation will complain about no trainable parameters found,
         // this is an extra step that TF's Python version does not need
@@ -325,33 +234,25 @@ export class MultiHeadAttention extends tf.layers.Layer {
             ...this.valueProjection.trainableWeights,
             ...this.outputProjection.trainableWeights
         ];
-
         // rename the weights otherwise they'll take on the default naming and overlap
         // each other which breaks model loading due to duplicate weight names
         let indexing = 0;
-
         for (const weight of this.trainableWeights) {
             const unique_name = `${this.getClassName()}_${indexing}`;
-            (weight as any).name += unique_name;
-            (weight as any).originalName += unique_name;
+            weight.name += unique_name;
+            weight.originalName += unique_name;
             indexing++;
         }
-
         super.build(inputShape);
     }
-
-
     /**
      * MultiHead attention's output is the same shape the query's.
      */
-    override computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[] {
+    computeOutputShape(inputShape) {
         return Array.isArray(inputShape) && Array.isArray(inputShape[0]) ? inputShape[0] : inputShape;
     }
-
-
-    override getConfig() {
+    getConfig() {
         const base_config = super.getConfig();
-
         const config = {
             numHeads: this.numHeads,
             embedDim: this.embedDim,
@@ -359,13 +260,10 @@ export class MultiHeadAttention extends tf.layers.Layer {
             causal: this.causal,
             dropout: this.dropout,
             name: this.name,
-        }
-
+        };
         Object.assign(config, base_config);
-
         return config;
     }
 }
-
-
 tf.serialization.registerClass(MultiHeadAttention);
+//# sourceMappingURL=multihead_attention.js.map

package/dist/layers/multihead_attention.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"multihead_attention.js","sourceRoot":"","sources":["../../src/layers/multihead_attention.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAGvC,OAAO,EAAE,MAAM,IAAI,kBAAkB,EAAE,MAAM,UAAU,CAAC;AAoBxD;;;;;;;;;;;;;;;;;;;;;;;;;;GA0BG;AACH,MAAM,OAAO,kBAAmB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACnD,MAAM,CAAC,SAAS,GAAG,oBAAoB,CAAC;IACrB,QAAQ,CAAS;IACjB,QAAQ,CAAS,CAAC,2DAA2D;IAC7E,OAAO,CAAU;IACjB,OAAO,CAAS;IAChB,MAAM,CAAU,CAAC,4CAA4C;IAEhF,mEAAmE;IACnE,wDAAwD;IACrC,eAAe,CAAkB;IACjC,aAAa,CAAkB;IAC/B,eAAe,CAAkB;IACjC,gBAAgB,CAAkB;IAGrD,YAAY,EAAE,QAAQ,EAAE,QAAQ,EAAE,OAAO,GAAG,IAAI,EAAE,OAAO,GAAG,GAAG,EAAE,MAAM,GAAG,KAAK,EAAE,GAAG,IAAI,EAA0B;QAC9G,KAAK,CAAC,IAAI,CAAC,CAAC;QAEZ,IAAI,QAAQ,GAAG,QAAQ,IAAI,CAAC,EAAE,CAAC;YAC3B,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,iBAAiB,IAAI,CAAC,IAAI,cAAc,QAAQ,mCAAmC,QAAQ,GAAG,CAAC,CAAC;QACtI,CAAC;QAED,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,OAAO,GAAG,OAAO,CAAC;QACvB,IAAI,CAAC,OAAO,GAAG,OAAO,CAAC;QACvB,IAAI,CAAC,MAAM,GAAG,MAAM,CAAC;QAErB,IAAI,IAAI,CAAC,OAAO,IAAI,CAAC,EAAE,CAAC;YACpB,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,6CAA6C,CAAC,CAAC;QACrF,CAAC;QAED,0DAA0D;QAC1D,gEAAgE;QAChE,IAAI,CAAC,eAAe,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,EAAE,OAAO,EAAE,KAAK,EAAE,QAAQ,EAAE,CAAC,CAAC;QACrE,IAAI,CAAC,aAAa,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,EAAE,OAAO,EAAE,KAAK,EAAE,QAAQ,EAAE,CAAC,CAAC;QACnE,IAAI,CAAC,eAAe,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,EAAE,OAAO,EAAE,KAAK,EAAE,QAAQ,EAAE,CAAC,CAAC;QACrE,IAAI,CAAC,gBAAgB,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,EAAE,OAAO,EAAE,KAAK,EAAE,QAAQ,EAAE,CAAC,CAAC;IAC1E,CAAC;IAGD;;;;;OAKG;IACM,IAAI,CACT,MAA+B,EAC/B,MAGC;QAED,6BAA6B;QAC7B,IAAI,CAAC,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,EAAE,CAAC;YACzB,MAAM,GAAG,CAAC,MAAM,CAAC,CAAC;QACtB,CAAC;QAED,6EAA6E;QAC7E,IAAI,MAAM,CAAC,MAAM,IAAI,CAAC,IAAI,MAAM,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YAC3C,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,UAAU,IAAI,CAAC,IAAI,gDAAgD,MAAM,CAAC,MAAM,gBAAgB,CAAC,CAAC;QACxI,CAAC;QAED,KAAK,MAAM,KAAK,IAAI,MAAM,EAAE,CAAC;YACzB,IAAI,KAAK,CAAC,KAAK,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;gBAC1B,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,UAAU,IAAI,CAAC,IAAI,0DAA0D,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;YAClJ,CAAC;QACL,CAAC;QAED,MAAM,CAAC,KAAK,EAAE,GAAG,EAAE,KAAK,CAAC,GAAG,MAAM,CAAC;QACnC,MAAM,WAAW,GAAG,MAAM,CAAC,WAAW,IAAI,IAAI,CAAC;QAC/C,MAAM,UAAU,GAAG,MAAM,CAAC,UAAU,IAAI,IAAI,CAAC;QAE7C,OAAO,MAAM,CAAC,MAAM,IAAI,CAAC;YACrB,kBAAkB;YAClB,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,KAAM,EAAE,GAAI,EAAE,KAAM,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,CAAC;YACrE,iBAAiB;YACjB,CAAC,CAAC,IAAI,CAAC,OAAO,CAAC,KAAM,EAAE,KAAM,EAAE,KAAM,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,CAAC,CAAC;IAChF,CAAC;IAGD;;OAEG;IACO,OAAO,CACb,WAAsB,EACtB,SAAoB,EACpB,WAAsB,EACtB,YAA8B,EAC9B,WAA6B,EAC7B,MAAc;QAEd,2BAA2B;QAC3B,+CAA+C;QAC/C,iDAAiD;QACjD,4CAA4C;QAC5C,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,EAAE,KAAK,EAAE,GAAG,EAAE,KAAK,EAAE,GAAG,IAAI,CAAC,qBAAqB,CAAC,WAAW,EAAE,SAAS,EAAE,WAAW,CAAC,CAAC;YAE9F,oGAAoG;YACpG,MAAM,qBAAqB,GAAG,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC;YAE3C,MAAM,EACF,WAAW,EAAE,SAAS,EAAE,WAAW,EACtC,GAAG,IAAI,CAAC,UAAU,CAAC,KAAK,EAAE,GAAG,EAAE,KAAK,EAAE,qBAAqB,CAAC,CAAC;YAE9D,gFAAgF;YAChF,MAAM,IAAI,GAAG,kBAAkB,CAAC,4BAA4B,CACxD,WAAW,EAAE,SAAS,EAAE,WAAW,EACnC,MAAM,CAAC,aAAa,IAAI,IAAI,EAAE,YAAY,EAAE,WAAW,EACvD,IAAI,CAAC,OAAO,EAAE,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;YAEvC,+CAA+C;YAC/C,MAAM,MAAM,GAAG,IAAI,CAAC,gBAAgB,CAAC,KAAK,CACtC,IAAI,CAAC,SAAS,CAAC,qBAAqB,CAAC,CAAC,OAAO,CAAC,CAAC,WAAW,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC,CAAC;YAE9F,OAAO,MAAmB,CAAC;QAC/B,CAAC,CAAC,CAAA;IACN,CAAC;IAGS,qBAAqB,CAAC,WAAsB,EAAE,SAAoB,EAAE,WAAsB;QAChG,wFAAwF;QACxF,2EAA2E;QAC3E,oHAAoH;QACpH,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,OAAO;gBACH,KAAK,EAAE,IAAI,CAAC,eAAe,CAAC,KAAK,CAAC,WAAW,CAAc;gBAC3D,GAAG,EAAE,IAAI,CAAC,aAAa,CAAC,KAAK,CAAC,SAAS,CAAc;gBACrD,KAAK,EAAE,IAAI,CAAC,eAAe,CAAC,KAAK,CAAC,WAAW,CAAc;aAC9D,CAAA;QACL,CAAC,CAAC,CAAA;IACN,CAAC;IAGS,UAAU,CAAC,KAAgB,EAAE,GAAc,EAAE,KAAgB,EAAE,OAAiB;QACtF,4EAA4E;QAC5E,2DAA2D;QAC3D,2FAA2F;QAC3F,MAAM,UAAU,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;QAClC,MAAM,WAAW,GAAG,CAAC,UAAU,EAAE,CAAC,CAAC,EAAE,IAAI,CAAC,QAAQ,EAAE,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,QAAQ,CAAC,CAAC;QAEnF,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,OAAO;gBACH,WAAW,EAAE,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,CAAC,SAAS,CAAC,OAAO,CAAgB;gBACzE,SAAS,EAAE,GAAG,CAAC,OAAO,CAAC,WAAW,CAAC,CAAC,SAAS,CAAC,OAAO,CAAgB;gBACrE,WAAW,EAAE,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,CAAC,SAAS,CAAC,OAAO,CAAgB;aAC5E,CAAA;QACL,CAAC,CAAC,CAAA;IACN,CAAC;IAGD;;;;;;;;;;OAUG;IACH,MAAM,CAAC,4BAA4B,CAC/B,KAAgB,EAChB,GAAc,EACd,KAAgB,EAChB,aAA+B,EAC/B,WAA6B,EAC7B,UAA4B,EAC5B,OAAe,EACf,MAAe,EACf,SAA6C,EAAE;QAE/C,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,EAAE,QAAQ,GAAG,KAAK,EAAE,cAAc,EAAE,GAAG,MAAM,CAAC;YAEpD,GAAG,CAAC,KAAK,CAAC,OAAO,CAAC,CAAC,GAAG,EAAE,KAAK,EAAE,EAAE;gBAC7B,IAAI,GAAG,CAAC,KAAK,CAAC,KAAK,CAAC,IAAI,KAAK,CAAC,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC;oBACzC,MAAM,KAAK,CAAC,sDAAsD;wBAC9D,gCAAgC,IAAI,CAAC,SAAS,CAAC,GAAG,CAAC,KAAK,CAAC,YAAY;wBACrE,IAAI,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC;gBACnD,CAAC;YACL,CAAC,CAAC,CAAA;YAGF,wFAAwF;YACxF,gEAAgE;YAChE,MAAM,iBAAiB,GAAG;gBACtB,KAAK,CAAC,KAAK,CAAC,KAAK,CAAC,KAAK,CAAC,MAAM,GAAG,CAAC,CAAC;gBACnC,GAAG,CAAC,KAAK,CAAC,GAAG,CAAC,KAAK,CAAC,MAAM,GAAG,CAAC,CAAC;aAAC,CAAC;YAErC,IAAI,IAAI,GAAG,EAAE,CAAC,KAAK,CAAC,iBAAiB,CAAC,CAAC;YAEvC,IAAI,MAAM,IAAI,iBAAiB,CAAC,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC;gBACrC,IAAI,aAAa,IAAI,aAAa,CAAC,KAAK,IAAI,MAAM,EAAE,CAAC;oBACjD,MAAM,KAAK,CAAC,oHAAoH,CAAC,CAAC;gBACtI,CAAC;gBAED,oFAAoF;gBACpF,2BAA2B;gBAC3B,IAAI,UAAU,EAAE,CAAC;oBACb,IAAI,GAAG,UAAU,CAAC;gBACtB,CAAC;qBAAM,CAAC;oBACJ,IAAI,GAAG,kBAAkB,CAAC,iBAAiB,CAAC,CAAC,CAAC,EAAE,iBAAiB,CAAC,CAAC,CAAC,CAAC,CAAC;gBAC1E,CAAC;YACL,CAAC;YAED,IAAI,aAAa,EAAE,CAAC;gBAChB,IAAI,aAAa,CAAC,KAAK,IAAI,MAAM,EAAE,CAAC;oBAChC,oCAAoC;oBACpC,4EAA4E;oBAC5E,IAAI,GAAG,IAAI,CAAC,GAAG,CAAC,aAAa,CAAC,IAAI,CAAC,SAAS,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC;gBACnE,CAAC;qBAAM,CAAC;oBACJ,sDAAsD;oBACtD,wFAAwF;oBACxF,IAAI,GAAG,aAAa,CAAC;gBACzB,CAAC;YACL,CAAC;YAED,8CAA8C;YAC9C,8BAA8B;YAC9B,iCAAiC;YACjC,wCAAwC;YACxC,mBAAmB;YACnB,+CAA+C;YAC/C,IAAI,WAAW,GAAG,KAAK;iBAClB,MAAM,CAAC,GAAG,EAAE,KAAK,EAAE,IAAI,CAAC;iBACxB,GAAG,CAAC,IAAI,CAAC,IAAI,CAAC,cAAc,IAAI,GAAG,CAAC,KAAK,CAAC,GAAG,CAAC,KAAK,CAAC,MAAM,GAAG,CAAC,CAAC,CAAC,CAAC;iBACjE,GAAG,CAAC,IAAI,CAAC,CAAC;YAEf,IAAI,WAAW,EAAE,CAAC;gBACd,sFAAsF;gBACtF,8BAA8B;gBAC9B,WAAW,GAAG,WAAW,CAAC,GAAG,CAAC,WAAW,CAAC,CAAC;YAC/C,CAAC;YAED,MAAM,IAAI,GAAG,EAAE,CAAC,OAAO,CAAC,WAAW,CAAC,CAAC;YAErC,MAAM,YAAY,GAAG,EAAE,CAAC,OAAO,CAAC,IAAI,EAAE,QAAQ,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YAC9D,MAAM,SAAS,GAAG,YAAY,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC;YAE7C,OAAO,SAAS,CAAC;QACrB,CAAC,CAAC,CAAC;IACP,CAAC;IAGQ,KAAK,CAAC,UAAiC;QAC5C,IAAI,WAAW,GAAe,EAAE,CAAC;QAEjC,IAAI,KAAK,CAAC,OAAO,CAAC,UAAU,CAAC,IAAI,KAAK,CAAC,OAAO,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC;YAC5D,WAAW,GAAG,UAAwB,CAAC;QAC3C,CAAC;aAAM,CAAC;YACJ,WAAW,GAAG,CAAC,UAAsB,EAAE,UAAsB,EAAE,UAAsB,CAAC,CAAC;QAC3F,CAAC;QAED,IAAI,WAAW,CAAC,MAAM,IAAI,CAAC,IAAI,WAAW,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YACrD,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,WAAW,IAAI,CAAC,IAAI,yDAAyD,IAAI,CAAC,SAAS,CAAC,UAAU,CAAC,EAAE,CAAC,CAAC;QACjJ,CAAC;QAED,kCAAkC;QAClC,IAAI,CAAC,eAAe,CAAC,KAAK,CAAC,WAAW,CAAC,CAAC,CAAC,CAAC,CAAC;QAC3C,IAAI,CAAC,aAAa,CAAC,KAAK,CAAC,WAAW,CAAC,CAAC,CAAC,CAAC,CAAC;QACzC,IAAI,CAAC,eAAe,CAAC,KAAK,CAAC,WAAW,CAAC,CAAC,CAAC,CAAC,CAAC;QAC3C,IAAI,CAAC,gBAAgB,CAAC,KAAK,CAAC,WAAW,CAAC,CAAC,CAAC,CAAC,CAAC;QAE5C,kEAAkE;QAClE,qEAAqE;QACrE,+DAA+D;QAC/D,IAAI,CAAC,gBAAgB,GAAG;YACpB,GAAG,IAAI,CAAC,eAAe,CAAC,gBAAgB;YACxC,GAAG,IAAI,CAAC,aAAa,CAAC,gBAAgB;YACtC,GAAG,IAAI,CAAC,eAAe,CAAC,gBAAgB;YACxC,GAAG,IAAI,CAAC,gBAAgB,CAAC,gBAAgB;SAC5C,CAAC;QAEF,8EAA8E;QAC9E,sEAAsE;QACtE,IAAI,QAAQ,GAAG,CAAC,CAAC;QAEjB,KAAK,MAAM,MAAM,IAAI,IAAI,CAAC,gBAAgB,EAAE,CAAC;YACzC,MAAM,WAAW,GAAG,GAAG,IAAI,CAAC,YAAY,EAAE,IAAI,QAAQ,EAAE,CAAC;YACxD,MAAc,CAAC,IAAI,IAAI,WAAW,CAAC;YACnC,MAAc,CAAC,YAAY,IAAI,WAAW,CAAC;YAC5C,QAAQ,EAAE,CAAC;QACf,CAAC;QAED,KAAK,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC;IAC5B,CAAC;IAGD;;OAEG;IACM,kBAAkB,CAAC,UAAiC;QACzD,OAAO,KAAK,CAAC,OAAO,CAAC,UAAU,CAAC,IAAI,KAAK,CAAC,OAAO,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,UAAU,CAAC;IAClG,CAAC;IAGQ,SAAS;QACd,MAAM,WAAW,GAAG,KAAK,CAAC,SAAS,EAAE,CAAC;QAEtC,MAAM,MAAM,GAAG;YACX,QAAQ,EAAE,IAAI,CAAC,QAAQ;YACvB,QAAQ,EAAE,IAAI,CAAC,QAAQ;YACvB,OAAO,EAAE,IAAI,CAAC,OAAO;YACrB,MAAM,EAAE,IAAI,CAAC,MAAM;YACnB,OAAO,EAAE,IAAI,CAAC,OAAO;YACrB,IAAI,EAAE,IAAI,CAAC,IAAI;SAClB,CAAA;QAED,MAAM,CAAC,MAAM,CAAC,MAAM,EAAE,WAAW,CAAC,CAAC;QAEnC,OAAO,MAAM,CAAC;IAClB,CAAC;;AAIL,EAAE,CAAC,aAAa,CAAC,aAAa,CAAC,kBAAkB,CAAC,CAAC"}

package/dist/layers/multihead_attention.test.d.ts

@@ -0,0 +1,2 @@
+export {};
+//# sourceMappingURL=multihead_attention.test.d.ts.map

package/dist/layers/multihead_attention.test.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"multihead_attention.test.d.ts","sourceRoot":"","sources":["../../src/layers/multihead_attention.test.ts"],"names":[],"mappings":""}