@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

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

@@ -1,100 +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)!,
+            numHeads: 2, embedDim: input.shape.at(-1),
             dropout: 0.5, activation: "gelu", dimsFeedForward: 321, useBias: false
         });
-
-        const output = decoder.apply(input) as tf.Tensor;
-
+        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 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)! });
+        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 });
+        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();
-    })
-
-
+        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");
-
+        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)! });
+        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)! });
+        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)! });
+        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)! });
+        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/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/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/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,+BAA+B,CAAC;AAGhG,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"}

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

@@ -1,23 +1,12 @@
 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 { MultiHeadAttention, type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
-
-
-export interface TransformerEncoderArgs extends MultiHeadAttentionArgs {
-    activation?: "relu" | "gelu";
-    dimsFeedForward?: number;
-}
-
-
+import { MultiHeadAttention } from "../layers/multihead_attention";
 /**
  * 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
@@ -28,28 +17,22 @@ export interface TransformerEncoderArgs extends MultiHeadAttentionArgs {
  */
 export class TransformerEncoder extends tf.layers.Layer {
     static className = "TransformerEncoder";
-
-    private readonly selfAttention: tf.layers.Layer;
-    private readonly selfAttentionDropout: tf.layers.Layer;
-    private readonly selfAttentionNorm: tf.layers.Layer;
-
-    private readonly reluLayer: tf.layers.Layer;
-    private readonly linearLayer: tf.layers.Layer;
-    private readonly feedForwardDropout: tf.layers.Layer;
-    private readonly feedFowardNorm: tf.layers.Layer;
-
-    private readonly numHeads: number;
-    private readonly embedDim: number;
-    private readonly causal: boolean;
-    private readonly useBias: boolean;
-    private readonly dropout: number;
-    private readonly activation: ActivationIdentifier;
-    private readonly dimsFeedForward: number;
-
-
-    constructor({ numHeads, embedDim, causal, useBias, dropout, activation, dimsFeedForward, ...args }: TransformerEncoderArgs) {
+    selfAttention;
+    selfAttentionDropout;
+    selfAttentionNorm;
+    reluLayer;
+    linearLayer;
+    feedForwardDropout;
+    feedFowardNorm;
+    numHeads;
+    embedDim;
+    causal;
+    useBias;
+    dropout;
+    activation;
+    dimsFeedForward;
+    constructor({ numHeads, embedDim, causal, useBias, dropout, activation, dimsFeedForward, ...args }) {
         super(args);
-
         this.numHeads = numHeads;
         this.embedDim = embedDim;
         this.causal = causal ?? false;
@@ -57,19 +40,16 @@ export class TransformerEncoder extends tf.layers.Layer {
         this.dropout = dropout ?? 0.1;
         this.activation = activation ?? "relu";
         this.dimsFeedForward = dimsFeedForward ?? 2048;
-
         if (this.dropout >= 1) {
             throw Error(`${this.getClassName()}::constructor dropout must be within [0, 1)`);
         }
-
         // self attention sub-block
         this.selfAttention = new MultiHeadAttention({
             numHeads: this.numHeads, embedDim: this.embedDim, useBias: this.useBias,
             dropout: this.dropout, causal: this.causal
         });
-        this.selfAttentionDropout = tf.layers.dropout({ rate: this.dropout })
+        this.selfAttentionDropout = tf.layers.dropout({ rate: this.dropout });
         this.selfAttentionNorm = tf.layers.layerNormalization({ epsilon: 1e-6 });
-
         // feed forward sub-block
         this.reluLayer = tf.layers.dense({
             units: this.dimsFeedForward, activation: this.activation,
@@ -82,96 +62,76 @@ export class TransformerEncoder extends tf.layers.Layer {
         this.feedForwardDropout = tf.layers.dropout({ rate: this.dropout });
         this.feedFowardNorm = tf.layers.layerNormalization({ epsilon: 1e-6 });
     }
-
-
     /**
      * Forward propagation
      */
-    override call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[] {
+    call(inputs, kwargs) {
         // validate the input tensors
-        let input: tf.Tensor;
-
+        let input;
         if (Array.isArray(inputs)) {
             if (inputs.length != 1) {
                 throw Error(`${this.getClassName}::call ${this.name} expects exactly 1 tensor` +
                     ` input, got ${inputs.length} inputs instead.`);
             }
-
             input = inputs[0];
-        } else {
+        }
+        else {
             input = inputs;
         }
-
         // perform forward propagation
         return tf.tidy(() => {
             const attention = this.selfAttentionBlock(input, kwargs);
             const feedforward = this.feedForwardBlock(attention, kwargs);
-
             return feedforward;
         });
     }
-
-
-    private selfAttentionBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor {
+    selfAttentionBlock(x, kwargs) {
         return tf.tidy(() => {
             const residual = x;
-
-            let attention = this.selfAttention.apply(x, kwargs) as tf.Tensor;
-            attention = this.selfAttentionDropout.apply(attention, kwargs) as tf.Tensor;
+            let attention = this.selfAttention.apply(x, kwargs);
+            attention = this.selfAttentionDropout.apply(attention, kwargs);
             attention = tf.add(attention, residual);
-            attention = this.selfAttentionNorm.apply(attention) as tf.Tensor;
-
+            attention = this.selfAttentionNorm.apply(attention);
             return attention;
         });
     }
-
-
-    private feedForwardBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor {
+    feedForwardBlock(x, kwargs) {
         return tf.tidy(() => {
             const residual = x;
-
             let feedForward = this.reluLayer.apply(x);
             feedForward = this.linearLayer.apply(feedForward);
-            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) as tf.Tensor;
-
+            feedForward = this.feedFowardNorm.apply(feedForward);
             return feedForward;
         });
     }
-
-
     /**
      * Initialize the sublayers' weights and track them to enable backpropagation.
      */
-    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];
         }
-
         // expects only 1 rank 3 tensor input
         if (input_shapes.length != 1 || input_shapes[0].length != 3) {
-            throw Error(`${this.getClassName()}::build ${this.name} expects a single input shape of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`)
+            throw Error(`${this.getClassName()}::build ${this.name} expects a single input shape of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`);
         }
-
         // initialize self attention sub-block's weights
         this.selfAttention.build(inputShape);
         this.selfAttentionNorm.build(inputShape);
-
         // inintialize feedforward sub-block's weights
         const reluLayerOutputShape = this.reluLayer.computeOutputShape(inputShape);
         const linearLayerOutputShape = this.linearLayer.computeOutputShape(reluLayerOutputShape);
-
         this.reluLayer.build(inputShape);
         this.linearLayer.build(reluLayerOutputShape);
         this.feedFowardNorm.build(linearLayerOutputShape);
-
         // track sublayers' weights
         this.trainableWeights = [
             ...this.selfAttention.trainableWeights,
@@ -182,28 +142,22 @@ export class TransformerEncoder 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(): tf.serialization.ConfigDict {
+    getConfig() {
         const base_config = super.getConfig();
-
         const config = {
             numHeads: this.numHeads,
             embedDim: this.embedDim,
@@ -213,12 +167,9 @@ export class TransformerEncoder extends tf.layers.Layer {
             activation: this.activation,
             dimsFeedForward: this.dimsFeedForward
         };
-
         Object.assign(config, base_config);
-
         return config;
     }
 }
-
-
 tf.serialization.registerClass(TransformerEncoder);
+//# sourceMappingURL=transformer_encoder.js.map

package/dist/layers/transformer_encoder.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"transformer_encoder.js","sourceRoot":"","sources":["../../src/layers/transformer_encoder.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAIvC,OAAO,EAAE,kBAAkB,EAA+B,MAAM,+BAA+B,CAAC;AAShG;;;;;;;;;;;;;;GAcG;AACH,MAAM,OAAO,kBAAmB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACnD,MAAM,CAAC,SAAS,GAAG,oBAAoB,CAAC;IAEvB,aAAa,CAAkB;IAC/B,oBAAoB,CAAkB;IACtC,iBAAiB,CAAkB;IAEnC,SAAS,CAAkB;IAC3B,WAAW,CAAkB;IAC7B,kBAAkB,CAAkB;IACpC,cAAc,CAAkB;IAEhC,QAAQ,CAAS;IACjB,QAAQ,CAAS;IACjB,MAAM,CAAU;IAChB,OAAO,CAAU;IACjB,OAAO,CAAS;IAChB,UAAU,CAAuB;IACjC,eAAe,CAAS;IAGzC,YAAY,EAAE,QAAQ,EAAE,QAAQ,EAAE,MAAM,EAAE,OAAO,EAAE,OAAO,EAAE,UAAU,EAAE,eAAe,EAAE,GAAG,IAAI,EAA0B;QACtH,KAAK,CAAC,IAAI,CAAC,CAAC;QAEZ,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,QAAQ,GAAG,QAAQ,CAAC;QACzB,IAAI,CAAC,MAAM,GAAG,MAAM,IAAI,KAAK,CAAC;QAC9B,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;QACvC,IAAI,CAAC,eAAe,GAAG,eAAe,IAAI,IAAI,CAAC;QAE/C,IAAI,IAAI,CAAC,OAAO,IAAI,CAAC,EAAE,CAAC;YACpB,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,6CAA6C,CAAC,CAAC;QACrF,CAAC;QAED,2BAA2B;QAC3B,IAAI,CAAC,aAAa,GAAG,IAAI,kBAAkB,CAAC;YACxC,QAAQ,EAAE,IAAI,CAAC,QAAQ,EAAE,QAAQ,EAAE,IAAI,CAAC,QAAQ,EAAE,OAAO,EAAE,IAAI,CAAC,OAAO;YACvE,OAAO,EAAE,IAAI,CAAC,OAAO,EAAE,MAAM,EAAE,IAAI,CAAC,MAAM;SAC7C,CAAC,CAAC;QACH,IAAI,CAAC,oBAAoB,GAAG,EAAE,CAAC,MAAM,CAAC,OAAO,CAAC,EAAE,IAAI,EAAE,IAAI,CAAC,OAAO,EAAE,CAAC,CAAA;QACrE,IAAI,CAAC,iBAAiB,GAAG,EAAE,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,OAAO,EAAE,IAAI,EAAE,CAAC,CAAC;QAEzE,yBAAyB;QACzB,IAAI,CAAC,SAAS,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;YAC7B,KAAK,EAAE,IAAI,CAAC,eAAe,EAAE,UAAU,EAAE,IAAI,CAAC,UAAU;YACxD,OAAO,EAAE,IAAI,CAAC,OAAO;SACxB,CAAC,CAAC;QACH,IAAI,CAAC,WAAW,GAAG,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;YAC/B,KAAK,EAAE,IAAI,CAAC,QAAQ,EAAE,UAAU,EAAE,QAAQ;YAC1C,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;;OAEG;IACM,IAAI,CAAC,MAA+B,EAAE,MAAc;QACzD,6BAA6B;QAC7B,IAAI,KAAgB,CAAC;QAErB,IAAI,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,EAAE,CAAC;YACxB,IAAI,MAAM,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;gBACrB,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,UAAU,IAAI,CAAC,IAAI,2BAA2B;oBAC1E,eAAe,MAAM,CAAC,MAAM,kBAAkB,CAAC,CAAC;YACxD,CAAC;YAED,KAAK,GAAG,MAAM,CAAC,CAAC,CAAC,CAAC;QACtB,CAAC;aAAM,CAAC;YACJ,KAAK,GAAG,MAAM,CAAC;QACnB,CAAC;QAED,8BAA8B;QAC9B,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,SAAS,GAAG,IAAI,CAAC,kBAAkB,CAAC,KAAK,EAAE,MAAM,CAAC,CAAC;YACzD,MAAM,WAAW,GAAG,IAAI,CAAC,gBAAgB,CAAC,SAAS,EAAE,MAAM,CAAC,CAAC;YAE7D,OAAO,WAAW,CAAC;QACvB,CAAC,CAAC,CAAC;IACP,CAAC;IAGO,kBAAkB,CAAC,CAAY,EAAE,MAAc;QACnD,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,QAAQ,GAAG,CAAC,CAAC;YAEnB,IAAI,SAAS,GAAG,IAAI,CAAC,aAAa,CAAC,KAAK,CAAC,CAAC,EAAE,MAAM,CAAc,CAAC;YACjE,SAAS,GAAG,IAAI,CAAC,oBAAoB,CAAC,KAAK,CAAC,SAAS,EAAE,MAAM,CAAc,CAAC;YAC5E,SAAS,GAAG,EAAE,CAAC,GAAG,CAAC,SAAS,EAAE,QAAQ,CAAC,CAAC;YACxC,SAAS,GAAG,IAAI,CAAC,iBAAiB,CAAC,KAAK,CAAC,SAAS,CAAc,CAAC;YAEjE,OAAO,SAAS,CAAC;QACrB,CAAC,CAAC,CAAC;IACP,CAAC;IAGO,gBAAgB,CAAC,CAAY,EAAE,MAAc;QACjD,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,QAAQ,GAAG,CAAC,CAAC;YAEnB,IAAI,WAAW,GAAG,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;YAC1C,WAAW,GAAG,IAAI,CAAC,WAAW,CAAC,KAAK,CAAC,WAAW,CAAC,CAAC;YAClD,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,CAAc,CAAC;YAElE,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,qCAAqC;QACrC,IAAI,YAAY,CAAC,MAAM,IAAI,CAAC,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC,MAAM,IAAI,CAAC,EAAE,CAAC;YAC1D,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,WAAW,IAAI,CAAC,IAAI,iEAAiE,IAAI,CAAC,SAAS,CAAC,UAAU,CAAC,EAAE,CAAC,CAAA;QACxJ,CAAC;QAED,gDAAgD;QAChD,IAAI,CAAC,aAAa,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC;QACrC,IAAI,CAAC,iBAAiB,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC;QAEzC,8CAA8C;QAC9C,MAAM,oBAAoB,GAAG,IAAI,CAAC,SAAS,CAAC,kBAAkB,CAAC,UAAU,CAAC,CAAC;QAC3E,MAAM,sBAAsB,GAAG,IAAI,CAAC,WAAW,CAAC,kBAAkB,CAAC,oBAAoB,CAAC,CAAC;QAEzF,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC;QACjC,IAAI,CAAC,WAAW,CAAC,KAAK,CAAC,oBAAoB,CAAC,CAAC;QAC7C,IAAI,CAAC,cAAc,CAAC,KAAK,CAAC,sBAAsB,CAAC,CAAC;QAElD,2BAA2B;QAC3B,IAAI,CAAC,gBAAgB,GAAG;YACpB,GAAG,IAAI,CAAC,aAAa,CAAC,gBAAgB;YACtC,GAAG,IAAI,CAAC,oBAAoB,CAAC,gBAAgB;YAC7C,GAAG,IAAI,CAAC,iBAAiB,CAAC,gBAAgB;YAC1C,GAAG,IAAI,CAAC,SAAS,CAAC,gBAAgB;YAClC,GAAG,IAAI,CAAC,WAAW,CAAC,gBAAgB;YACpC,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,MAAM,EAAE,IAAI,CAAC,MAAM;YACnB,OAAO,EAAE,IAAI,CAAC,OAAO;YACrB,OAAO,EAAE,IAAI,CAAC,OAAO;YACrB,UAAU,EAAE,IAAI,CAAC,UAAU;YAC3B,eAAe,EAAE,IAAI,CAAC,eAAe;SACxC,CAAC;QAEF,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/transformer_encoder.test.d.ts

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

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

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

package/{src/layers/transformer_encoder.test.ts → dist/layers/transformer_encoder.test.js}

@@ -1,85 +1,58 @@
 import * as tf from "@tensorflow/tfjs";
-
 import { TransformerEncoder } from "@/layers/transformer_encoder";
-
 // disables warning for using the faster node backend,
 // https://github.com/tensorflow/tfjs/issues/5349#issuecomment-885170504
 tf.env().set('IS_NODE', false);
-
-
 describe("TransformerEncoder tests", () => {
     it("should return an output with the same shape as the input", () => {
         const input = tf.randomUniform([2, 3, 10]);
-
         const decoder = new TransformerEncoder({
-            numHeads: 2, embedDim: input.shape.at(-1)!,
+            numHeads: 2, embedDim: input.shape.at(-1),
             dropout: 0.5, activation: "gelu", dimsFeedForward: 512, useBias: true
         });
-
-        const output = decoder.apply(input) as tf.Tensor;
-
+        const output = decoder.apply(input);
         expect(output.shape.length).toBe(input.shape.length);
-    })
-
-
+    });
     test("correct forward calls", () => {
         const input = tf.randomUniform([2, 3, 10]);
-
-        const encoder = new TransformerEncoder({ numHeads: 2, embedDim: input.shape.at(-1)! });
+        const encoder = new TransformerEncoder({ numHeads: 2, embedDim: input.shape.at(-1) });
         expect(() => encoder.apply(input)).not.toThrow();
         expect(() => encoder.apply([input])).not.toThrow();
-
-        const causal = new TransformerEncoder({ numHeads: 2, embedDim: input.shape.at(-1)!, causal: true });
+        const causal = new TransformerEncoder({ 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, 10]);
-
-        expect(() => new TransformerEncoder({ numHeads: 3, embedDim: input.shape.at(-1)! })).toThrow();
-        expect(() => new TransformerEncoder({ numHeads: 5, embedDim: input.shape.at(-1)! })).not.toThrow();
-    })
-
-
+        expect(() => new TransformerEncoder({ numHeads: 3, embedDim: input.shape.at(-1) })).toThrow();
+        expect(() => new TransformerEncoder({ numHeads: 5, embedDim: input.shape.at(-1) })).not.toThrow();
+    });
     it("should not accept non-rank 3 tensor inputs", () => {
         const incorrect_input = tf.randomUniform([2, 3, 10, 10]);
         const incorrect_input2 = tf.randomUniform([2, 3]);
         const correct_input = tf.randomUniform([2, 3, 10]);
-
-
-        const encoder = new TransformerEncoder({ numHeads: 2, embedDim: incorrect_input.shape.at(-1)! });
+        const encoder = new TransformerEncoder({ numHeads: 2, embedDim: incorrect_input.shape.at(-1) });
         expect(() => encoder.apply([correct_input, correct_input])).toThrow();
-
         expect(() => encoder.apply(incorrect_input)).toThrow();
         expect(() => encoder.apply(incorrect_input2)).toThrow();
-
         expect(() => encoder.apply([correct_input, incorrect_input])).toThrow();
         expect(() => encoder.apply([incorrect_input, correct_input])).toThrow();
-
         expect(() => encoder.apply([correct_input, incorrect_input2])).toThrow();
         expect(() => encoder.apply([incorrect_input2, correct_input])).toThrow();
-    })
-
-
+    });
     it("should accept exactly one input", () => {
         const input = tf.randomUniform([2, 3, 10]);
-
-        const encoder = new TransformerEncoder({ numHeads: 1, embedDim: input.shape.at(-1)! });
+        const encoder = new TransformerEncoder({ numHeads: 1, embedDim: input.shape.at(-1) });
         expect(() => encoder.apply(input)).not.toThrow();
         expect(() => encoder.apply([input])).not.toThrow();
-
         expect(() => encoder.apply([])).toThrow();
         expect(() => encoder.apply([input, input])).toThrow();
-        expect(() => encoder.apply([input, input, input])).toThrow()
-    })
-
-
+        expect(() => encoder.apply([input, input, input])).toThrow();
+    });
     it("should return a non-empty config dict", () => {
         const input = tf.randomUniform([2, 3, 10]);
-
-        const encoder = new TransformerEncoder({ numHeads: 1, embedDim: input.shape.at(-1)! });
+        const encoder = new TransformerEncoder({ numHeads: 1, embedDim: input.shape.at(-1) });
         expect(Object.keys(encoder.getConfig())).not.toBe(0);
-    })
-})
+    });
+});
+//# sourceMappingURL=transformer_encoder.test.js.map

package/dist/layers/transformer_encoder.test.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"transformer_encoder.test.js","sourceRoot":"","sources":["../../src/layers/transformer_encoder.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,IAAI;SACxE,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,uBAAuB,EAAE,GAAG,EAAE;QAC/B,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,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,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,OAAO,EAAE,CAAC;QAC/F,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;IACvG,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,4CAA4C,EAAE,GAAG,EAAE;QAClD,MAAM,eAAe,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC;QACzD,MAAM,gBAAgB,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAClD,MAAM,aAAa,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAGnD,MAAM,OAAO,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,eAAe,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACjG,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,aAAa,EAAE,aAAa,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAEtE,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,eAAe,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACvD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,gBAAgB,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAExD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,aAAa,EAAE,eAAe,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACxE,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,eAAe,EAAE,aAAa,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAExE,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,aAAa,EAAE,gBAAgB,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACzE,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,gBAAgB,EAAE,aAAa,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IAC7E,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,iCAAiC,EAAE,GAAG,EAAE;QACvC,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,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,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAC1C,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACtD,MAAM,CAAC,GAAG,EAAE,CAAC,OAAO,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAA;IAChE,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/losses/dice.d.ts

@@ -0,0 +1,30 @@
+import * as tf from "@tensorflow/tfjs";
+export declare function diceBinaryStandard(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor;
+export declare function diceBinaryGlobal(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor;
+export declare function diceCategoricalStandard(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor;
+export declare function diceCategoricalGeneralized(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor;
+export declare function diceCategoricalGlobal(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor;
+/**
+ * Calculates the Sorensen-Dice coefficient and the binary cross entropy losses.
+ * Both have equal weight.
+ *
+ * @param y_true the label tensor
+ * @param y_pred the prediction tensor (not sparse)
+ * @returns a tensor of shape `[ batch ]`
+ */
+export declare function diceBinaryCrossentropy(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor;
+/**
+ * Calculates the Sorensen-Dice coefficient and the categorical cross entropy losses.
+ * Both have equal weight. Expects dense (non-sparse) label tensors.
+ *
+ * This does not support sparse tensors because TFJS's
+ * sparseCategoricalCrossentropy loss onehots the label
+ * and calls categoricalCrossentropy. See
+ * https://github.com/tensorflow/tfjs/blob/0fc04d958ea592f3b8db79a8b3b497b5c8904097/tfjs-layers/src/losses.ts#L143-L146
+ *
+ * @param y_true the label
+ * @param y_pred the prediction tensor (not sparse)
+ * @returns a tensor of shape `[ batch ]`
+ */
+export declare function diceCategoricalCrossentropy(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor;
+//# sourceMappingURL=dice.d.ts.map

package/dist/losses/dice.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"dice.d.ts","sourceRoot":"","sources":["../../src/losses/dice.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAWvC,wBAAgB,kBAAkB,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,CAclF;AAQD,wBAAgB,gBAAgB,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,CAahF;AAOD,wBAAgB,uBAAuB,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,CAUvF;AAOD,wBAAgB,0BAA0B,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,CAgB1F;AAOD,wBAAgB,qBAAqB,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,CAWrF;AAOD;;;;;;;GAOG;AACH,wBAAgB,sBAAsB,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,CAMtF;AAOD;;;;;;;;;;;;GAYG;AACH,wBAAgB,2BAA2B,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,CAM3F"}