@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

package/dist/src/layers/token_and_positional_embedding.test.js

@@ -0,0 +1,58 @@
+import * as tf from '@tensorflow/tfjs';
+import { TokenAndPositionalEmbedding } from '@/layers/token_and_positional_embedding';
+// disables warning for using the faster node backend,
+// https://github.com/tensorflow/tfjs/issues/5349#issuecomment-885170504
+tf.env().set('IS_NODE', false);
+describe("PositionalEncoding tests", () => {
+    test("layer initialization", () => {
+        expect(() => new TokenAndPositionalEmbedding({ maxSequenceLength: 0, embedDim: 10, vocabularySize: 10_000 })).toThrow();
+        expect(() => new TokenAndPositionalEmbedding({ embedDim: 0, vocabularySize: 10_000 })).toThrow();
+        expect(() => new TokenAndPositionalEmbedding({ embedDim: 10, vocabularySize: 0 })).toThrow();
+        expect(() => new TokenAndPositionalEmbedding({ embedDim: 10, vocabularySize: 10_000 })).not.toThrow();
+        expect(() => new TokenAndPositionalEmbedding({ embedDim: 10, vocabularySize: 10_000 })).not.toThrow();
+    });
+    test("successfull forward calls", () => {
+        const embed_dims = 32;
+        const sequences = 4;
+        const vocab_size = 10_000;
+        const input = tf.randomUniform([2, sequences]);
+        const embedding = new TokenAndPositionalEmbedding({ embedDim: embed_dims, dropout: 0.1, vocabularySize: vocab_size });
+        expect(() => embedding.apply(input)).not.toThrow();
+        expect(() => embedding.apply([input])).not.toThrow();
+    });
+    test("layer build", () => {
+        const input_ok = tf.randomUniform([2, 4]);
+        const input_too_many_words = tf.randomUniform([2, 700]);
+        const input_is_image = tf.randomUniform([1, 32, 32, 3]);
+        let embedding = new TokenAndPositionalEmbedding({ embedDim: 32, maxSequenceLength: 500, vocabularySize: 1_000 });
+        expect(() => embedding.build(input_ok.shape)).not.toThrow();
+        embedding = new TokenAndPositionalEmbedding({ embedDim: 32, maxSequenceLength: 500, vocabularySize: 1_000 });
+        expect(() => embedding.build([input_ok.shape, input_ok.shape])).not.toThrow();
+        new TokenAndPositionalEmbedding({ embedDim: 32, maxSequenceLength: 500, vocabularySize: 1_000 });
+        expect(() => embedding.build(input_too_many_words.shape)).toThrow();
+        expect(() => embedding.build(input_is_image.shape)).toThrow();
+    });
+    it("should throw when more than one input provided, input sequences are too large, or incorrect input rank", () => {
+        const sequences_too_long = tf.randomUniform([10, 1000]);
+        const multiple_correct_inputs = [tf.randomUniform([2, 3]), tf.randomUniform([2, 3])];
+        const wrong_rank = tf.randomUniform([10, 32, 32]);
+        const positional = new TokenAndPositionalEmbedding({ maxSequenceLength: 10, embedDim: 32, vocabularySize: 10_000 });
+        positional.build([2, 3]); // get past the initial build call to test forward prop
+        expect(() => positional.apply(sequences_too_long)).toThrow();
+        expect(() => positional.apply(multiple_correct_inputs)).toThrow();
+        expect(() => positional.apply(wrong_rank)).toThrow();
+    });
+    it("should return a non-empty config dict", () => {
+        const embedding = new TokenAndPositionalEmbedding({ embedDim: 32, vocabularySize: 10_000 });
+        expect(Object.keys(embedding.getConfig())).not.toBe(0);
+    });
+    it("should return an output shape of [batch, sequences, embed dims]", () => {
+        const words = 100;
+        const batch = 2;
+        const embed_dims = 64;
+        const input = tf.randomUniform([batch, words]);
+        const embedding = new TokenAndPositionalEmbedding({ embedDim: embed_dims, vocabularySize: 10_000 });
+        expect(embedding.computeOutputShape(input.shape)).toEqual([batch, words, embed_dims]);
+    });
+});
+//# sourceMappingURL=token_and_positional_embedding.test.js.map

package/dist/src/layers/token_and_positional_embedding.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"token_and_positional_embedding.test.js","sourceRoot":"","sources":["../../../src/layers/token_and_positional_embedding.test.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAEvC,OAAO,EAAE,2BAA2B,EAAE,MAAM,yCAAyC,CAAC;AAEtF,sDAAsD;AACtD,wEAAwE;AACxE,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,CAAC,SAAS,EAAE,KAAK,CAAC,CAAC;AAG/B,QAAQ,CAAC,0BAA0B,EAAE,GAAG,EAAE;IACtC,IAAI,CAAC,sBAAsB,EAAE,GAAG,EAAE;QAC9B,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,2BAA2B,CAAC,EAAE,iBAAiB,EAAE,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,cAAc,EAAE,MAAM,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACxH,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,cAAc,EAAE,MAAM,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACjG,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,cAAc,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAE7F,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,cAAc,EAAE,MAAM,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACtG,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,cAAc,EAAE,MAAM,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;IAC1G,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,2BAA2B,EAAE,GAAG,EAAE;QACnC,MAAM,UAAU,GAAG,EAAE,CAAC;QACtB,MAAM,SAAS,GAAG,CAAC,CAAC;QACpB,MAAM,UAAU,GAAG,MAAM,CAAC;QAC1B,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC,CAAC;QAE/C,MAAM,SAAS,GAAG,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,UAAU,EAAE,OAAO,EAAE,GAAG,EAAE,cAAc,EAAE,UAAU,EAAE,CAAC,CAAC;QACtH,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACnD,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;IACzD,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,aAAa,EAAE,GAAG,EAAE;QACrB,MAAM,QAAQ,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAC1C,MAAM,oBAAoB,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC,CAAC;QACxD,MAAM,cAAc,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,CAAC;QAExD,IAAI,SAAS,GAAG,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,iBAAiB,EAAE,GAAG,EAAE,cAAc,EAAE,KAAK,EAAE,CAAC,CAAC;QACjH,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,QAAQ,CAAC,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAE5D,SAAS,GAAG,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,iBAAiB,EAAE,GAAG,EAAE,cAAc,EAAE,KAAK,EAAE,CAAC,CAAC;QAC7G,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,QAAQ,CAAC,KAAK,EAAE,QAAQ,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAE9E,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,iBAAiB,EAAE,GAAG,EAAE,cAAc,EAAE,KAAK,EAAE,CAAC,CAAC;QACjG,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,oBAAoB,CAAC,KAAK,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACpE,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,cAAc,CAAC,KAAK,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IAClE,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,wGAAwG,EAAE,GAAG,EAAE;QAC9G,MAAM,kBAAkB,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,EAAE,EAAE,IAAI,CAAC,CAAC,CAAC;QACxD,MAAM,uBAAuB,GAAG,CAAC,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;QACrF,MAAM,UAAU,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC;QAElD,MAAM,UAAU,GAAG,IAAI,2BAA2B,CAAC,EAAE,iBAAiB,EAAE,EAAE,EAAE,QAAQ,EAAE,EAAE,EAAE,cAAc,EAAE,MAAM,EAAE,CAAC,CAAC;QACpH,UAAU,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,uDAAuD;QAEjF,MAAM,CAAC,GAAG,EAAE,CAAC,UAAU,CAAC,KAAK,CAAC,kBAAkB,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAC7D,MAAM,CAAC,GAAG,EAAE,CAAC,UAAU,CAAC,KAAK,CAAC,uBAAuB,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAClE,MAAM,CAAC,GAAG,EAAE,CAAC,UAAU,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IACzD,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,uCAAuC,EAAE,GAAG,EAAE;QAC7C,MAAM,SAAS,GAAG,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,cAAc,EAAE,MAAM,EAAE,CAAC,CAAC;QAC5F,MAAM,CAAC,MAAM,CAAC,IAAI,CAAC,SAAS,CAAC,SAAS,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;IAC3D,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,iEAAiE,EAAE,GAAG,EAAE;QACvE,MAAM,KAAK,GAAG,GAAG,CAAC;QAClB,MAAM,KAAK,GAAG,CAAC,CAAC;QAChB,MAAM,UAAU,GAAG,EAAE,CAAC;QAEtB,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC;QAE/C,MAAM,SAAS,GAAG,IAAI,2BAA2B,CAAC,EAAE,QAAQ,EAAE,UAAU,EAAE,cAAc,EAAE,MAAM,EAAE,CAAC,CAAC;QAEpG,MAAM,CAAC,SAAS,CAAC,kBAAkB,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,UAAU,CAAC,CAAC,CAAC;IAC1F,CAAC,CAAC,CAAA;AACN,CAAC,CAAC,CAAC"}

package/dist/src/layers/transformer_decoder.d.ts

@@ -0,0 +1,69 @@
+import * as tf from "@tensorflow/tfjs";
+import { type Kwargs } from "@tensorflow/tfjs-layers/dist/types";
+import { type ActivationIdentifier } from "@tensorflow/tfjs-layers/dist/keras_format/activation_config";
+import { type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
+export interface TransformerDecoderArgs extends Omit<MultiHeadAttentionArgs, "causal"> {
+    activation?: "relu" | "gelu";
+    dimsFeedForward?: number;
+    causal?: boolean;
+}
+/**
+ * This class implements the transformer decoder architecture from
+ * the 2017 paper "Attention Is All You Need".
+ *
+ * This decoder-only transformer layer accepts one tensor input.
+ * The input tensor should have the shape
+ * `[ batch, sequences, embedding dims ]`.
+ *
+ * Causal masking is enabled by default for the initial attention sub-layer.
+ *
+ * @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 on inputs (masks future inputs to prevent looking ahead), default `true`
+ * @param dropout use dropout during the attention calculations, default `0.1`
+ * @param activation the activation of the intermediate feed forward layer, default `relu`
+ * @param dimsFeedForward the size of the intermediate feed forward layer, default `2048`
+ * @param useBias use bias for the dense sublayers and multiHead attention's dense sublayers, default `true`
+ */
+export declare class TransformerDecoder extends tf.layers.Layer {
+    static className: string;
+    protected readonly causalSelfAttention: tf.layers.Layer;
+    protected readonly causalSelfAttentionDropout: tf.layers.Layer;
+    protected readonly causalSelfAttentionNorm: tf.layers.Layer;
+    protected readonly feedforward1: tf.layers.Layer;
+    protected readonly feedforward2: tf.layers.Layer;
+    protected readonly feedForwardDropout: tf.layers.Layer;
+    protected readonly feedFowardNorm: tf.layers.Layer;
+    protected readonly numHeads: number;
+    protected readonly embedDim: number;
+    protected readonly useBias: boolean;
+    protected readonly dropout: number;
+    protected readonly activation: ActivationIdentifier;
+    protected readonly dimsFeedForward: number;
+    constructor({ numHeads, embedDim, useBias, dropout, activation, dimsFeedForward, ...args }: TransformerDecoderArgs);
+    /**
+     * Forward propagation
+     *
+     * @param inputs input tensor
+     * @return the output tensor
+     */
+    call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[];
+    protected causalSelfAttentionBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor;
+    protected feedForwardBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor;
+    /**
+     * Initialize the sublayers' weights and track them to enable serialization
+     */
+    build(inputShape: tf.Shape | tf.Shape[]): void;
+    /**
+     * Save the layer's hyperparameters for serialization
+     */
+    getConfig(): {
+        numHeads: number;
+        embedDim: number;
+        useBias: boolean;
+        dropout: number;
+        activation: ActivationIdentifier;
+        dimsFeedForward: number;
+    };
+}
+//# sourceMappingURL=transformer_decoder.d.ts.map

package/dist/src/layers/transformer_decoder.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"transformer_decoder.d.ts","sourceRoot":"","sources":["../../../src/layers/transformer_decoder.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,KAAK,MAAM,EAAE,MAAM,oCAAoC,CAAC;AACjE,OAAO,EAAE,KAAK,oBAAoB,EAAE,MAAM,6DAA6D,CAAC;AAExG,OAAO,EAAE,KAAK,sBAAsB,EAAE,MAAM,8BAA8B,CAAC;AAI3E,MAAM,WAAW,sBAAuB,SAAQ,IAAI,CAAC,sBAAsB,EAAE,QAAQ,CAAC;IAClF,UAAU,CAAC,EAAE,MAAM,GAAG,MAAM,CAAC;IAC7B,eAAe,CAAC,EAAE,MAAM,CAAC;IACzB,MAAM,CAAC,EAAE,OAAO,CAAC;CACpB;AAGD;;;;;;;;;;;;;;;;;GAiBG;AACH,qBAAa,kBAAmB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACnD,MAAM,CAAC,SAAS,SAAwB;IAExC,SAAS,CAAC,QAAQ,CAAC,mBAAmB,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IACxD,SAAS,CAAC,QAAQ,CAAC,0BAA0B,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IAC/D,SAAS,CAAC,QAAQ,CAAC,uBAAuB,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IAE5D,SAAS,CAAC,QAAQ,CAAC,YAAY,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IACjD,SAAS,CAAC,QAAQ,CAAC,YAAY,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IACjD,SAAS,CAAC,QAAQ,CAAC,kBAAkB,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IACvD,SAAS,CAAC,QAAQ,CAAC,cAAc,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IAEnD,SAAS,CAAC,QAAQ,CAAC,QAAQ,EAAE,MAAM,CAAC;IACpC,SAAS,CAAC,QAAQ,CAAC,QAAQ,EAAE,MAAM,CAAC;IACpC,SAAS,CAAC,QAAQ,CAAC,OAAO,EAAE,OAAO,CAAC;IACpC,SAAS,CAAC,QAAQ,CAAC,OAAO,EAAE,MAAM,CAAC;IACnC,SAAS,CAAC,QAAQ,CAAC,UAAU,EAAE,oBAAoB,CAAC;IACpD,SAAS,CAAC,QAAQ,CAAC,eAAe,EAAE,MAAM,CAAC;gBAE/B,EAAE,QAAQ,EAAE,QAAQ,EAAE,OAAO,EAAE,OAAO,EAAE,UAAU,EAAE,eAAe,EAAE,GAAG,IAAI,EAAE,EAAE,sBAAsB;IAyClH;;;;;OAKG;IACM,IAAI,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE,EAAE,MAAM,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE;IAoBvF,SAAS,CAAC,wBAAwB,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM;IAc3E,SAAS,CAAC,gBAAgB,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM;IAenE;;OAEG;IACM,KAAK,CAAC,UAAU,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,IAAI;IA6DvD;;OAEG;IACM,SAAS;;;;;;;;CAiBrB"}

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

@@ -1,28 +1,15 @@
 import * as tf from "@tensorflow/tfjs";
-import { type Kwargs } from "@tensorflow/tfjs-layers/dist/types";
-import { type ActivationIdentifier } from "@tensorflow/tfjs-layers/dist/keras_format/activation_config";
-
-import { type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
 import { CachedRoPEMultiHeadAttention } from "@/layers/cached_rope_multihead_attention";
-
-
-export interface TransformerDecoderArgs extends Omit<MultiHeadAttentionArgs, "causal"> {
-    activation?: "relu" | "gelu";
-    dimsFeedForward?: number;
-    causal?: boolean; // use causal mask for attention on inputs
-}
-
-
 /**
  * This class implements the transformer decoder architecture from
  * the 2017 paper "Attention Is All You Need".
- * 
+ *
  * This decoder-only transformer layer accepts one tensor input.
  * The input tensor should have the shape
  * `[ batch, sequences, embedding dims ]`.
- * 
+ *
  * Causal masking is enabled by default for the initial attention sub-layer.
- * 
+ *
  * @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 on inputs (masks future inputs to prevent looking ahead), default `true`
@@ -33,48 +20,39 @@ export interface TransformerDecoderArgs extends Omit<MultiHeadAttentionArgs, "ca
  */
 export class TransformerDecoder extends tf.layers.Layer {
     static className = "TransformerDecoder";
-
-    protected readonly causalSelfAttention: tf.layers.Layer;
-    protected readonly causalSelfAttentionDropout: tf.layers.Layer;
-    protected readonly causalSelfAttentionNorm: tf.layers.Layer;
-
-    protected readonly feedforward1: tf.layers.Layer;
-    protected readonly feedforward2: tf.layers.Layer;
-    protected readonly feedForwardDropout: tf.layers.Layer;
-    protected readonly feedFowardNorm: tf.layers.Layer;
-
-    protected readonly numHeads: number;
-    protected readonly embedDim: number;
-    protected readonly useBias: boolean;
-    protected readonly dropout: number;
-    protected readonly activation: ActivationIdentifier;
-    protected readonly dimsFeedForward: number;
-
-    constructor({ numHeads, embedDim, useBias, dropout, activation, dimsFeedForward, ...args }: TransformerDecoderArgs) {
+    causalSelfAttention;
+    causalSelfAttentionDropout;
+    causalSelfAttentionNorm;
+    feedforward1;
+    feedforward2;
+    feedForwardDropout;
+    feedFowardNorm;
+    numHeads;
+    embedDim;
+    useBias;
+    dropout;
+    activation;
+    dimsFeedForward;
+    constructor({ numHeads, embedDim, useBias, dropout, activation, dimsFeedForward, ...args }) {
         super(args);
-
         this.numHeads = numHeads;
         this.embedDim = embedDim;
         this.useBias = useBias ?? true;
         this.dropout = dropout ?? 0.1;
         this.activation = activation ?? "relu";
-
         if (this.dropout >= 1) {
             throw Error(`${this.getClassName()}::constructor dropout must be within [0, 1)`);
         }
-
         // in the paper section 3.3, d_model=512 (embedDim) and first dense layer outputs d_ff=2048
         this.dimsFeedForward = dimsFeedForward ?? embedDim * 4;
-
         // self attention sub-block
         this.causalSelfAttention = new CachedRoPEMultiHeadAttention({
             numHeads: this.numHeads, embedDim: this.embedDim,
             useBias: this.useBias, dropout: this.dropout,
             causal: true
         });
-        this.causalSelfAttentionDropout = tf.layers.dropout({ rate: this.dropout })
+        this.causalSelfAttentionDropout = tf.layers.dropout({ rate: this.dropout });
         this.causalSelfAttentionNorm = tf.layers.layerNormalization({ epsilon: 1e-6 });
-
         // feed forward sub-block
         this.feedforward1 = tf.layers.dense({
             units: this.dimsFeedForward,
@@ -89,101 +67,79 @@ export class TransformerDecoder extends tf.layers.Layer {
         this.feedForwardDropout = tf.layers.dropout({ rate: this.dropout });
         this.feedFowardNorm = tf.layers.layerNormalization({ epsilon: 1e-6 });
     }
-
-
     /**
      * Forward propagation
-     * 
+     *
      * @param inputs input tensor
      * @return the output tensor
      */
-    override call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[] {
+    call(inputs, kwargs) {
         // validate the input tensors
         if (Array.isArray(inputs) && inputs.length != 1 && inputs.length != 2) {
             throw Error(`${this.getClassName()}::call ${this.name} expects one input tensor, got ${inputs.length} inputs.`);
         }
-
         if (Array.isArray(inputs)) {
-            inputs = inputs[0] as tf.Tensor;
+            inputs = inputs[0];
         }
-
         // perform forward propagation
         return tf.tidy(() => {
             let output = this.causalSelfAttentionBlock(inputs, kwargs);
             output = this.feedForwardBlock(output, kwargs);
-
             return output;
         });
     }
-
-
-    protected causalSelfAttentionBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor {
+    causalSelfAttentionBlock(x, kwargs) {
         return tf.tidy(() => {
             const residual = x;
-
-            let attention = this.causalSelfAttention.apply(x, kwargs) as tf.Tensor;
-            attention = this.causalSelfAttentionDropout.apply(attention, kwargs) as tf.Tensor;
+            let attention = this.causalSelfAttention.apply(x, kwargs);
+            attention = this.causalSelfAttentionDropout.apply(attention, kwargs);
             attention = tf.add(attention, residual);
-            attention = this.causalSelfAttentionNorm.apply(attention, kwargs) as tf.Tensor;
-
+            attention = this.causalSelfAttentionNorm.apply(attention, kwargs);
             return attention;
         });
     }
-
-
-    protected feedForwardBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor {
+    feedForwardBlock(x, kwargs) {
         return tf.tidy(() => {
             const residual = x;
-
             let feedForward = this.feedforward1.apply(x, kwargs);
             feedForward = this.feedforward2.apply(feedForward, kwargs);
-            feedForward = this.feedForwardDropout.apply(feedForward, kwargs) as tf.Tensor;
+            feedForward = this.feedForwardDropout.apply(feedForward, kwargs);
             feedForward = tf.add(feedForward, residual);
-            feedForward = this.feedFowardNorm.apply(feedForward, kwargs) as tf.Tensor;
-
+            feedForward = this.feedFowardNorm.apply(feedForward, kwargs);
             return feedForward;
         });
     }
-
-
     /**
      * Initialize the sublayers' weights and track them to enable serialization
      */
-    override build(inputShape: tf.Shape | tf.Shape[]): void {
-        let input_shapes: tf.Shape[] = [];
-
+    build(inputShape) {
+        let input_shapes = [];
         if (Array.isArray(inputShape) && Array.isArray(inputShape[0])) {
             // input is an array of shapes
-            input_shapes = inputShape as tf.Shape[];
-        } else if (inputShape.length != 0) {
+            input_shapes = inputShape;
+        }
+        else if (inputShape.length != 0) {
             // input is a single shape
-            input_shapes = [inputShape as tf.Shape];
+            input_shapes = [inputShape];
         }
-
         if (input_shapes.length != 1 && input_shapes.length != 2) {
             throw Error(`${this.getClassName()}::build ${this.name} expects an input shape` +
-                ` of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`)
+                ` of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`);
         }
-
         const [decoderInputShape] = input_shapes;
-
         if (decoderInputShape?.length != 3) {
             throw Error(`${this.getClassName()}::build ${this.name} expects an input shape` +
-                ` of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`)
+                ` of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`);
         }
-
         // initialize causal self attention sub-block's weights
         this.causalSelfAttention.build(decoderInputShape);
         this.causalSelfAttentionNorm.build(this.causalSelfAttention.computeOutputShape(decoderInputShape));
-
         // initialize feedforward sub-block's weights
         const feedforward1OutputShape = this.feedforward1.computeOutputShape(decoderInputShape);
         const feedforward2OutputShape = this.feedforward2.computeOutputShape(feedforward1OutputShape);
-
         this.feedforward1.build(decoderInputShape);
         this.feedforward2.build(feedforward1OutputShape);
         this.feedFowardNorm.build(feedforward2OutputShape);
-
         // track sublayers' weights
         this.trainableWeights = [
             ...this.causalSelfAttention.trainableWeights,
@@ -194,28 +150,22 @@ export class TransformerDecoder extends tf.layers.Layer {
             ...this.feedForwardDropout.trainableWeights,
             ...this.feedFowardNorm.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);
     }
-
-
     /**
      * Save the layer's hyperparameters for serialization
      */
-    override getConfig() {
+    getConfig() {
         const base_config = super.getConfig();
-
         const config = {
             numHeads: this.numHeads,
             embedDim: this.embedDim,
@@ -223,14 +173,10 @@ export class TransformerDecoder extends tf.layers.Layer {
             dropout: this.dropout,
             activation: this.activation,
             dimsFeedForward: this.dimsFeedForward
-        }
-
+        };
         Object.assign(config, base_config);
-
         return config;
     }
-
 }
-
-
 tf.serialization.registerClass(TransformerDecoder);
+//# sourceMappingURL=transformer_decoder.js.map

package/dist/src/layers/transformer_decoder.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"transformer_decoder.js","sourceRoot":"","sources":["../../../src/layers/transformer_decoder.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAKvC,OAAO,EAAE,4BAA4B,EAAE,MAAM,0CAA0C,CAAC;AAUxF;;;;;;;;;;;;;;;;;GAiBG;AACH,MAAM,OAAO,kBAAmB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACnD,MAAM,CAAC,SAAS,GAAG,oBAAoB,CAAC;IAErB,mBAAmB,CAAkB;IACrC,0BAA0B,CAAkB;IAC5C,uBAAuB,CAAkB;IAEzC,YAAY,CAAkB;IAC9B,YAAY,CAAkB;IAC9B,kBAAkB,CAAkB;IACpC,cAAc,CAAkB;IAEhC,QAAQ,CAAS;IACjB,QAAQ,CAAS;IACjB,OAAO,CAAU;IACjB,OAAO,CAAS;IAChB,UAAU,CAAuB;IACjC,eAAe,CAAS;IAE3C,YAAY,EAAE,QAAQ,EAAE,QAAQ,EAAE,OAAO,EAAE,OAAO,EAAE,UAAU,EAAE,eAAe,EAAE,GAAG,IAAI,EAA0B;QAC9G,KAAK,CAAC,IAAI,CAAC,CAAC;QAEZ,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,OAAO,GAAG,OAAO,IAAI,IAAI,CAAC;QAC/B,IAAI,CAAC,OAAO,GAAG,OAAO,IAAI,GAAG,CAAC;QAC9B,IAAI,CAAC,UAAU,GAAG,UAAU,IAAI,MAAM,CAAC;QAEvC,IAAI,IAAI,CAAC,OAAO,IAAI,CAAC,EAAE,CAAC;YACpB,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,6CAA6C,CAAC,CAAC;QACrF,CAAC;QAED,2FAA2F;QAC3F,IAAI,CAAC,eAAe,GAAG,eAAe,IAAI,QAAQ,GAAG,CAAC,CAAC;QAEvD,2BAA2B;QAC3B,IAAI,CAAC,mBAAmB,GAAG,IAAI,4BAA4B,CAAC;YACxD,QAAQ,EAAE,IAAI,CAAC,QAAQ,EAAE,QAAQ,EAAE,IAAI,CAAC,QAAQ;YAChD,OAAO,EAAE,IAAI,CAAC,OAAO,EAAE,OAAO,EAAE,IAAI,CAAC,OAAO;YAC5C,MAAM,EAAE,IAAI;SACf,CAAC,CAAC;QACH,IAAI,CAAC,0BAA0B,GAAG,EAAE,CAAC,MAAM,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,IAAI,CAAC,OAAO,EAAE,CAAC,CAAA;QAC3E,IAAI,CAAC,uBAAuB,GAAG,EAAE,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,OAAO,EAAE,IAAI,EAAE,CAAC,CAAC;QAE/E,yBAAyB;QACzB,IAAI,CAAC,YAAY,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;YAChC,KAAK,EAAE,IAAI,CAAC,eAAe;YAC3B,UAAU,EAAE,IAAI,CAAC,UAAU;YAC3B,OAAO,EAAE,IAAI,CAAC,OAAO;SACxB,CAAC,CAAC;QACH,IAAI,CAAC,YAAY,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;YAChC,KAAK,EAAE,IAAI,CAAC,QAAQ;YACpB,UAAU,EAAE,QAAQ;YACpB,OAAO,EAAE,IAAI,CAAC,OAAO;SACxB,CAAC,CAAC;QACH,IAAI,CAAC,kBAAkB,GAAG,EAAE,CAAC,MAAM,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,IAAI,CAAC,OAAO,EAAE,CAAC,CAAC;QACpE,IAAI,CAAC,cAAc,GAAG,EAAE,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,OAAO,EAAE,IAAI,EAAE,CAAC,CAAC;IAC1E,CAAC;IAGD;;;;;OAKG;IACM,IAAI,CAAC,MAA+B,EAAE,MAAc;QACzD,6BAA6B;QAC7B,IAAI,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,IAAI,MAAM,CAAC,MAAM,IAAI,CAAC,IAAI,MAAM,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YACpE,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,UAAU,IAAI,CAAC,IAAI,kCAAkC,MAAM,CAAC,MAAM,UAAU,CAAC,CAAC;QACpH,CAAC;QAED,IAAI,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,EAAE,CAAC;YACxB,MAAM,GAAG,MAAM,CAAC,CAAC,CAAc,CAAC;QACpC,CAAC;QAED,8BAA8B;QAC9B,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,IAAI,MAAM,GAAG,IAAI,CAAC,wBAAwB,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;YAC3D,MAAM,GAAG,IAAI,CAAC,gBAAgB,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;YAE/C,OAAO,MAAM,CAAC;QAClB,CAAC,CAAC,CAAC;IACP,CAAC;IAGS,wBAAwB,CAAC,CAAY,EAAE,MAAc;QAC3D,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,QAAQ,GAAG,CAAC,CAAC;YAEnB,IAAI,SAAS,GAAG,IAAI,CAAC,mBAAmB,CAAC,KAAK,CAAC,CAAC,EAAE,MAAM,CAAc,CAAC;YACvE,SAAS,GAAG,IAAI,CAAC,0BAA0B,CAAC,KAAK,CAAC,SAAS,EAAE,MAAM,CAAc,CAAC;YAClF,SAAS,GAAG,EAAE,CAAC,GAAG,CAAC,SAAS,EAAE,QAAQ,CAAC,CAAC;YACxC,SAAS,GAAG,IAAI,CAAC,uBAAuB,CAAC,KAAK,CAAC,SAAS,EAAE,MAAM,CAAc,CAAC;YAE/E,OAAO,SAAS,CAAC;QACrB,CAAC,CAAC,CAAC;IACP,CAAC;IAGS,gBAAgB,CAAC,CAAY,EAAE,MAAc;QACnD,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,QAAQ,GAAG,CAAC,CAAC;YAEnB,IAAI,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,KAAK,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YACrD,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,KAAK,CAAC,WAAW,EAAE,MAAM,CAAC,CAAC;YAC3D,WAAW,GAAG,IAAI,CAAC,kBAAkB,CAAC,KAAK,CAAC,WAAW,EAAE,MAAM,CAAc,CAAC;YAC9E,WAAW,GAAG,EAAE,CAAC,GAAG,CAAC,WAAW,EAAE,QAAQ,CAAC,CAAC;YAC5C,WAAW,GAAG,IAAI,CAAC,cAAc,CAAC,KAAK,CAAC,WAAW,EAAE,MAAM,CAAc,CAAC;YAE1E,OAAO,WAAW,CAAC;QACvB,CAAC,CAAC,CAAC;IACP,CAAC;IAGD;;OAEG;IACM,KAAK,CAAC,UAAiC;QAC5C,IAAI,YAAY,GAAe,EAAE,CAAC;QAElC,IAAI,KAAK,CAAC,OAAO,CAAC,UAAU,CAAC,IAAI,KAAK,CAAC,OAAO,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC,EAAE,CAAC;YAC5D,8BAA8B;YAC9B,YAAY,GAAG,UAAwB,CAAC;QAC5C,CAAC;aAAM,IAAI,UAAU,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YAChC,0BAA0B;YAC1B,YAAY,GAAG,CAAC,UAAsB,CAAC,CAAC;QAC5C,CAAC;QAED,IAAI,YAAY,CAAC,MAAM,IAAI,CAAC,IAAI,YAAY,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YACvD,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,WAAW,IAAI,CAAC,IAAI,yBAAyB;gBAC3E,oCAAoC,IAAI,CAAC,SAAS,CAAC,UAAU,CAAC,EAAE,CAAC,CAAA;QACzE,CAAC;QAED,MAAM,CAAC,iBAAiB,CAAC,GAAG,YAAY,CAAC;QAEzC,IAAI,iBAAiB,EAAE,MAAM,IAAI,CAAC,EAAE,CAAC;YACjC,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,WAAW,IAAI,CAAC,IAAI,yBAAyB;gBAC3E,oCAAoC,IAAI,CAAC,SAAS,CAAC,UAAU,CAAC,EAAE,CAAC,CAAA;QACzE,CAAC;QAED,uDAAuD;QACvD,IAAI,CAAC,mBAAmB,CAAC,KAAK,CAAC,iBAAiB,CAAC,CAAC;QAClD,IAAI,CAAC,uBAAuB,CAAC,KAAK,CAAC,IAAI,CAAC,mBAAmB,CAAC,kBAAkB,CAAC,iBAAiB,CAAC,CAAC,CAAC;QAEnG,6CAA6C;QAC7C,MAAM,uBAAuB,GAAG,IAAI,CAAC,YAAY,CAAC,kBAAkB,CAAC,iBAAiB,CAAC,CAAC;QACxF,MAAM,uBAAuB,GAAG,IAAI,CAAC,YAAY,CAAC,kBAAkB,CAAC,uBAAuB,CAAC,CAAC;QAE9F,IAAI,CAAC,YAAY,CAAC,KAAK,CAAC,iBAAiB,CAAC,CAAC;QAC3C,IAAI,CAAC,YAAY,CAAC,KAAK,CAAC,uBAAuB,CAAC,CAAC;QACjD,IAAI,CAAC,cAAc,CAAC,KAAK,CAAC,uBAAuB,CAAC,CAAC;QAEnD,2BAA2B;QAC3B,IAAI,CAAC,gBAAgB,GAAG;YACpB,GAAG,IAAI,CAAC,mBAAmB,CAAC,gBAAgB;YAC5C,GAAG,IAAI,CAAC,0BAA0B,CAAC,gBAAgB;YACnD,GAAG,IAAI,CAAC,uBAAuB,CAAC,gBAAgB;YAChD,GAAG,IAAI,CAAC,YAAY,CAAC,gBAAgB;YACrC,GAAG,IAAI,CAAC,YAAY,CAAC,gBAAgB;YACrC,GAAG,IAAI,CAAC,kBAAkB,CAAC,gBAAgB;YAC3C,GAAG,IAAI,CAAC,cAAc,CAAC,gBAAgB;SAC1C,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,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,OAAO,EAAE,IAAI,CAAC,OAAO;YACrB,UAAU,EAAE,IAAI,CAAC,UAAU;YAC3B,eAAe,EAAE,IAAI,CAAC,eAAe;SACxC,CAAA;QAED,MAAM,CAAC,MAAM,CAAC,MAAM,EAAE,WAAW,CAAC,CAAC;QAEnC,OAAO,MAAM,CAAC;IAClB,CAAC;;AAKL,EAAE,CAAC,aAAa,CAAC,aAAa,CAAC,kBAAkB,CAAC,CAAC"}

package/dist/src/layers/transformer_decoder.test.d.ts

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

package/dist/src/layers/transformer_decoder.test.d.ts.map

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

package/dist/src/layers/transformer_decoder.test.js

@@ -0,0 +1,72 @@
+import * as tf from '@tensorflow/tfjs';
+import { TransformerDecoder } from '@/layers/transformer_decoder';
+// disables warning for using the faster node backend,
+// https://github.com/tensorflow/tfjs/issues/5349#issuecomment-885170504
+tf.env().set('IS_NODE', false);
+describe("TransformerDecoder tests", () => {
+    it("should return an output with the same shape as the input", () => {
+        const input = tf.randomUniform([2, 3, 12]);
+        const decoder = new TransformerDecoder({
+            numHeads: 2, embedDim: input.shape.at(-1),
+            dropout: 0.5, activation: "gelu", dimsFeedForward: 321, useBias: false
+        });
+        const output = decoder.apply(input);
+        expect(output.shape.length).toBe(input.shape.length);
+    });
+    test("forward calls", () => {
+        const input = tf.randomUniform([2, 3, 12]);
+        const mask = tf.randomUniform([input.shape[0], input.shape[1]], -1, 2, "bool");
+        const incorrect_mask = tf.randomUniform([2, 5, 12], -1, 2, "bool");
+        const decoder = new TransformerDecoder({ numHeads: 2, embedDim: input.shape.at(-1) });
+        expect(() => decoder.apply(input)).not.toThrow();
+        expect(() => decoder.apply([input])).not.toThrow();
+        // causal masking
+        const causal = new TransformerDecoder({ numHeads: 2, embedDim: input.shape.at(-1), causal: true });
+        expect(() => causal.apply(input)).not.toThrow();
+        expect(() => causal.apply([input])).not.toThrow();
+    });
+    it("should fail to instantiate a layer if heads count is not divisible by the input's embedding dimension", () => {
+        const input = tf.randomUniform([2, 3, 12]);
+        expect(() => new TransformerDecoder({ numHeads: 3, embedDim: input.shape.at(-1) })).not.toThrow();
+        expect(() => new TransformerDecoder({ numHeads: 5, embedDim: input.shape.at(-1) })).toThrow();
+    });
+    it("should not accept non-rank 3 tensor inputs", () => {
+        const embed_dim = 12;
+        const BAD_RANK4 = tf.randomUniform([2, 3, 12, embed_dim]);
+        const BAD_RANK2 = tf.randomUniform([2, embed_dim]);
+        const GOOD = tf.randomUniform([2, 3, embed_dim]);
+        const mask = tf.randomUniform([GOOD.shape[0], GOOD.shape[1]], -1, 2, "bool");
+        let decoder = new TransformerDecoder({ numHeads: 2, embedDim: embed_dim });
+        // BAD
+        expect(() => decoder.apply(BAD_RANK4)).toThrow();
+        expect(() => decoder.apply(BAD_RANK2)).toThrow();
+        // OK
+        decoder = new TransformerDecoder({ numHeads: 2, embedDim: embed_dim });
+        expect(() => decoder.apply(GOOD)).not.toThrow();
+        expect(() => decoder.apply([GOOD])).not.toThrow();
+        expect(() => decoder.apply([GOOD, mask])).not.toThrow();
+    });
+    it("should not accept inputs that are less or more than 1 and 2 tensors", () => {
+        const input = tf.randomUniform([2, 3, 12]);
+        let decoder = new TransformerDecoder({ numHeads: 1, embedDim: input.shape.at(-1) });
+        // OK
+        expect(() => decoder.apply(input)).not.toThrow();
+        expect(() => decoder.apply([input])).not.toThrow();
+        // BAD
+        decoder = new TransformerDecoder({ numHeads: 1, embedDim: input.shape.at(-1) });
+        expect(() => decoder.apply([])).toThrow(); // stops at build()
+        decoder.apply(input); // get past the initial build
+        expect(() => decoder.apply([input, input, input])).toThrow();
+        expect(() => decoder.apply([input, input, input, input])).toThrow();
+        // BAD (tests build())
+        decoder = new TransformerDecoder({ numHeads: 1, embedDim: input.shape.at(-1) });
+        expect(() => decoder.apply([input, input, input])).toThrow();
+        expect(() => decoder.apply([input, input, input, input])).toThrow();
+    });
+    it("should return a non-empty config dict", () => {
+        const input = tf.randomUniform([2, 3, 12]);
+        const decoder = new TransformerDecoder({ numHeads: 1, embedDim: input.shape.at(-1) });
+        expect(Object.keys(decoder.getConfig())).not.toBe(0);
+    });
+});
+//# sourceMappingURL=transformer_decoder.test.js.map

package/dist/src/layers/transformer_decoder.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"transformer_decoder.test.js","sourceRoot":"","sources":["../../../src/layers/transformer_decoder.test.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAEvC,OAAO,EAAE,kBAAkB,EAAE,MAAM,8BAA8B,CAAC;AAElE,sDAAsD;AACtD,wEAAwE;AACxE,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,CAAC,SAAS,EAAE,KAAK,CAAC,CAAC;AAG/B,QAAQ,CAAC,0BAA0B,EAAE,GAAG,EAAE;IACtC,EAAE,CAAC,0DAA0D,EAAE,GAAG,EAAE;QAChE,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAE3C,MAAM,OAAO,GAAG,IAAI,kBAAkB,CAAC;YACnC,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE;YAC1C,OAAO,EAAE,GAAG,EAAE,UAAU,EAAE,MAAM,EAAE,eAAe,EAAE,GAAG,EAAE,OAAO,EAAE,KAAK;SACzE,CAAC,CAAC;QAEH,MAAM,MAAM,GAAG,OAAO,CAAC,KAAK,CAAC,KAAK,CAAc,CAAC;QAEjD,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,MAAM,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,MAAM,CAAC,CAAC;IACzD,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,eAAe,EAAE,GAAG,EAAE;QACvB,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAC3C,MAAM,IAAI,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAE,EAAE,KAAK,CAAC,KAAK,CAAC,CAAC,CAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,MAAM,CAAC,CAAC;QACjF,MAAM,cAAc,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,MAAM,CAAC,CAAC;QAGnE,MAAM,OAAO,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACvF,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACjD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAEnD,iBAAiB;QACjB,MAAM,MAAM,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,MAAM,EAAE,IAAI,EAAE,CAAC,CAAC;QACpG,MAAM,CAAC,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAChD,MAAM,CAAC,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;IACtD,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,uGAAuG,EAAE,GAAG,EAAE;QAC7G,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAE3C,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACnG,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IACnG,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,4CAA4C,EAAE,GAAG,EAAE;QAClD,MAAM,SAAS,GAAG,EAAE,CAAC;QAErB,MAAM,SAAS,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,SAAS,CAAC,CAAC,CAAC;QAC1D,MAAM,SAAS,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,SAAS,CAAC,CAAC,CAAC;QACnD,MAAM,IAAI,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,SAAS,CAAC,CAAC,CAAC;QACjD,MAAM,IAAI,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,CAAE,EAAE,IAAI,CAAC,KAAK,CAAC,CAAC,CAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,MAAM,CAAC,CAAC;QAE/E,IAAI,OAAO,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,SAAS,EAAE,CAAC,CAAC;QAE3E,MAAM;QACN,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,SAAS,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACjD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,SAAS,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAEjD,KAAK;QACL,OAAO,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,SAAS,EAAE,CAAC,CAAC;QACvE,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,IAAI,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAChD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAClD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,IAAI,EAAE,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;IAC5D,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,qEAAqE,EAAE,GAAG,EAAE;QAC3E,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAE3C,IAAI,OAAO,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACrF,KAAK;QACL,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACjD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAEnD,MAAM;QACN,OAAO,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACjF,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC,CAAC,mBAAmB;QAC9D,OAAO,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,6BAA6B;QACnD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAC7D,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAEpE,sBAAsB;QACtB,OAAO,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACjF,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAC7D,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IACxE,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,uCAAuC,EAAE,GAAG,EAAE;QAC7C,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAE3C,MAAM,OAAO,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACvF,MAAM,CAAC,MAAM,CAAC,IAAI,CAAC,OAAO,CAAC,SAAS,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;IACzD,CAAC,CAAC,CAAA;AACN,CAAC,CAAC,CAAA"}

package/dist/src/layers/transformer_encoder.d.ts

@@ -0,0 +1,55 @@
+import * as tf from "@tensorflow/tfjs";
+import { type Kwargs } from "@tensorflow/tfjs-layers/dist/types";
+import { type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
+export interface TransformerEncoderArgs extends MultiHeadAttentionArgs {
+    activation?: "relu" | "gelu";
+    dimsFeedForward?: number;
+}
+/**
+ * This class implements the transformer encoder architecture from the 2017 paper
+ * Attention Is All You Need.
+ *
+ * This layer accepts exactly one tensor input with the shape
+ * `[ batch, sequences, embedding dims ]`.
+ *
+ * @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` for encoders
+ * @param dropout use dropout during the attention calculations, default `0.1`
+ * @param activation the activation of the intermediate feed forward layer, default `relu`
+ * @param dimsFeedForward the size of the intermediate feed forward layer, default `2048`
+ * @param useBias use bias for the dense sublayers and multiHead attention's dense sublayers, default `true`
+ */
+export declare class TransformerEncoder extends tf.layers.Layer {
+    static className: string;
+    private readonly selfAttention;
+    private readonly selfAttentionDropout;
+    private readonly selfAttentionNorm;
+    private readonly reluLayer;
+    private readonly linearLayer;
+    private readonly feedForwardDropout;
+    private readonly feedFowardNorm;
+    private readonly numHeads;
+    private readonly embedDim;
+    private readonly causal;
+    private readonly useBias;
+    private readonly dropout;
+    private readonly activation;
+    private readonly dimsFeedForward;
+    constructor({ numHeads, embedDim, causal, useBias, dropout, activation, dimsFeedForward, ...args }: TransformerEncoderArgs);
+    /**
+     * Forward propagation
+     */
+    call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[];
+    private selfAttentionBlock;
+    private feedForwardBlock;
+    /**
+     * Initialize the sublayers' weights and track them to enable backpropagation.
+     */
+    build(inputShape: tf.Shape | tf.Shape[]): void;
+    /**
+     * Save the layer's hyperparameters for serialization
+     */
+    getConfig(): tf.serialization.ConfigDict;
+}
+//# sourceMappingURL=transformer_encoder.d.ts.map

package/dist/src/layers/transformer_encoder.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"transformer_encoder.d.ts","sourceRoot":"","sources":["../../../src/layers/transformer_encoder.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,KAAK,MAAM,EAAE,MAAM,oCAAoC,CAAC;AAGjE,OAAO,EAAsB,KAAK,sBAAsB,EAAE,MAAM,8BAA8B,CAAC;AAG/F,MAAM,WAAW,sBAAuB,SAAQ,sBAAsB;IAClE,UAAU,CAAC,EAAE,MAAM,GAAG,MAAM,CAAC;IAC7B,eAAe,CAAC,EAAE,MAAM,CAAC;CAC5B;AAGD;;;;;;;;;;;;;;GAcG;AACH,qBAAa,kBAAmB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACnD,MAAM,CAAC,SAAS,SAAwB;IAExC,OAAO,CAAC,QAAQ,CAAC,aAAa,CAAkB;IAChD,OAAO,CAAC,QAAQ,CAAC,oBAAoB,CAAkB;IACvD,OAAO,CAAC,QAAQ,CAAC,iBAAiB,CAAkB;IAEpD,OAAO,CAAC,QAAQ,CAAC,SAAS,CAAkB;IAC5C,OAAO,CAAC,QAAQ,CAAC,WAAW,CAAkB;IAC9C,OAAO,CAAC,QAAQ,CAAC,kBAAkB,CAAkB;IACrD,OAAO,CAAC,QAAQ,CAAC,cAAc,CAAkB;IAEjD,OAAO,CAAC,QAAQ,CAAC,QAAQ,CAAS;IAClC,OAAO,CAAC,QAAQ,CAAC,QAAQ,CAAS;IAClC,OAAO,CAAC,QAAQ,CAAC,MAAM,CAAU;IACjC,OAAO,CAAC,QAAQ,CAAC,OAAO,CAAU;IAClC,OAAO,CAAC,QAAQ,CAAC,OAAO,CAAS;IACjC,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAuB;IAClD,OAAO,CAAC,QAAQ,CAAC,eAAe,CAAS;gBAG7B,EAAE,QAAQ,EAAE,QAAQ,EAAE,MAAM,EAAE,OAAO,EAAE,OAAO,EAAE,UAAU,EAAE,eAAe,EAAE,GAAG,IAAI,EAAE,EAAE,sBAAsB;IAqC1H;;OAEG;IACM,IAAI,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE,EAAE,MAAM,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE;IAyBvF,OAAO,CAAC,kBAAkB;IAc1B,OAAO,CAAC,gBAAgB;IAexB;;OAEG;IACM,KAAK,CAAC,UAAU,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,IAAI;IAsDvD;;OAEG;IACM,SAAS,IAAI,EAAE,CAAC,aAAa,CAAC,UAAU;CAiBpD"}