@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

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

@@ -1,212 +1,160 @@
 import * as tf from '@tensorflow/tfjs';
-
 import { CachedRoPEMultiHeadAttention } from '@/layers/cached_rope_multihead_attention';
-import { generateCausalAttentionMask } from '@/utils';
+import { causal as generateCausalMask } from "@/masks";
 import { MultiHeadAttention } from '@/layers/multihead_attention';
-
 // disables warning for using the faster node backend,
 // https://github.com/tensorflow/tfjs/issues/5349#issuecomment-885170504
 tf.env().set('IS_NODE', false);
-
-
 describe("MultiHeadAttention tests", () => {
     it("should fail to instantiate a layer if heads count is not divisible by the input's embedding dimension", () => {
         expect(() => new CachedRoPEMultiHeadAttention({ numHeads: 3, embedDim: 10 })).toThrow();
         expect(() => new CachedRoPEMultiHeadAttention({ numHeads: 15, embedDim: 60 })).not.toThrow();
-    })
-
-
+    });
     test("successfull forward calls", () => {
         const input = tf.randomUniform([2, 3, 12]);
-
-        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1)! });
+        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1) });
         expect(() => attention.apply(input)).not.toThrow();
         expect(() => attention.apply([input])).not.toThrow();
-
-        const causal = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1)!, causal: true });
+        const causal = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1), causal: true });
         expect(() => causal.apply(input)).not.toThrow();
         expect(() => causal.apply([input])).not.toThrow();
-    })
-
-
+    });
     test("query and value must have the same shape for scaled dot product attention to succeed", () => {
         const query = tf.randomUniform([2, 3, 12]);
         const key = tf.randomUniform([2, 3, 12]);
         const value = tf.randomUniform([2, 3, 12]);
         const value_thats_too_long = tf.randomUniform([2, 100, 12]);
-
-        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: query.shape.at(-1)! });
+        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: query.shape.at(-1) });
         expect(() => attention.apply([query, key, value])).not.toThrow();
         expect(() => attention.apply([query, key, value_thats_too_long])).toThrow();
-    })
-
-
+    });
     it("should only accept rank 3 tensors", () => {
         const embed_dims = 12;
-
         const BAD_RANK2 = tf.randomUniform([2, embed_dims]);
         const GOOD = tf.randomUniform([2, 3, embed_dims]);
         const BAD_RANK4 = tf.randomUniform([2, 3, 10, embed_dims]);
-
         const attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: embed_dims });
-
         // BAD
         expect(() => attention.apply(BAD_RANK2)).toThrow();
         expect(() => attention.apply([BAD_RANK2])).toThrow();
         expect(() => attention.apply([BAD_RANK2, BAD_RANK2, BAD_RANK2])).toThrow();
-
         // OK
         expect(() => attention.apply(GOOD)).not.toThrow();
         expect(() => attention.apply([GOOD])).not.toThrow();
         expect(() => attention.apply([GOOD, GOOD, GOOD])).not.toThrow();
-
         // BAD
         expect(() => attention.apply(BAD_RANK4)).toThrow();
         expect(() => attention.apply([BAD_RANK4])).toThrow();
         expect(() => attention.apply([BAD_RANK4, BAD_RANK4, BAD_RANK4])).toThrow();
-
         // BAD
         expect(() => attention.apply([GOOD, BAD_RANK2, BAD_RANK4])).toThrow();
         expect(() => attention.apply([BAD_RANK2, GOOD, BAD_RANK4])).toThrow();
         expect(() => attention.apply([BAD_RANK2, BAD_RANK4, GOOD])).toThrow();
         expect(() => attention.apply([BAD_RANK2, GOOD, GOOD])).toThrow();
         expect(() => attention.apply([GOOD, GOOD, BAD_RANK4])).toThrow();
-    })
-
-
+    });
     it("should only 1 or 3 inputs total", () => {
         const input = tf.randomUniform([2, 3, 12]);
-
-        let attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1)! });
-
+        let attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1) });
         // OK
         expect(() => attention.apply(input, { packingMask: undefined })).not.toThrow();
         expect(() => attention.apply([input])).not.toThrow();
         // reinitialize to rerun build()
-        attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1)! });
+        attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1) });
         expect(() => attention.apply([input, input, input])).not.toThrow();
-
         // BAD
         expect(() => attention.apply([])).toThrow();
         expect(() => attention.apply([input, input])).toThrow();
         // reinitialize to rerun build()
-        attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1)! });
+        attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: input.shape.at(-1) });
         expect(() => attention.apply([input, input, input, input])).toThrow();
-    })
-
-
+    });
     test("attention masking", () => {
         const query = tf.randomUniform([2, 3, 12]);
         const key = tf.randomUniform([2, 3, 12]);
         const value = tf.randomUniform([2, 3, 12]);
-
-        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: query.shape.at(-1)!, causal: true });
-
+        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 2, embedDim: query.shape.at(-1), causal: true });
         expect(() => attention.call(query, {})).not.toThrow();
-
         // cross attention
         expect(() => attention.call([query, key, value], {})).not.toThrow();
-
-
         const query5 = tf.randomUniform([2, 5, 10]);
         const key4 = tf.randomUniform([2, 4, 10]);
         const value5 = tf.randomUniform([2, 4, 10]);
-
         const expected_mask = tf.tensor([[
-            // vertical represents query, false means that token cannot attend to the keys
-            // horizontal represents key, false means that token cannot attend to the queries
-            [false, false, false, false],
-            [true, true, true, false,],
-            [true, true, true, false,],
-            [false, false, false, false],
-            [true, true, true, false,],
-        ]]);
-
+                // vertical represents query, false means that token cannot attend to the keys
+                // horizontal represents key, false means that token cannot attend to the queries
+                [false, false, false, false],
+                [true, true, true, false,],
+                [true, true, true, false,],
+                [false, false, false, false],
+                [true, true, true, false,],
+            ]]);
         const packing_mask = tf.tensor([
             [0, 0, 0, -1e7, -1e7],
             [0, 0, 0, -1e7, -1e7],
             [0, 0, 0, -1e7, -1e7],
             [-1e7, -1e7, -1e7, 0, 0],
             [-1e7, -1e7, -1e7, 0, 0]
-        ])
-
+        ]);
         // for causal attention, the attention mask must be boolean
         expect(() => MultiHeadAttention.scaledDotProductionAttention(query5, key4, value5, expected_mask.asType("float32"), null, null, 0.1, true, { scaling_factor: 10 })).toThrow();
         // for causal attention, using pre-calculated causal mask
-        expect(() => MultiHeadAttention.scaledDotProductionAttention(query5, key4, value5, expected_mask.asType("float32"), null, generateCausalAttentionMask(query5.shape[1]!, key4.shape[1]!), 0.2, true, { scaling_factor: 10 })).toThrow();
+        expect(() => MultiHeadAttention.scaledDotProductionAttention(query5, key4, value5, expected_mask.asType("float32"), null, generateCausalMask(query5.shape[1], key4.shape[1]), 0.2, true, { scaling_factor: 10 })).toThrow();
         // when not using causal attention, the attention mask can be a float32 tensor
         expect(() => MultiHeadAttention.scaledDotProductionAttention(query5, key4, value5, expected_mask.asType("float32"), null, null, 0, false)).not.toThrow();
         // packing mask for self attention
         expect(() => MultiHeadAttention.scaledDotProductionAttention(query5, query5, query5, null, packing_mask, null, 0.9, true)).not.toThrow();
-    })
-
-
+    });
     it("should return a non-empty config dict", () => {
         const input = tf.randomUniform([2, 3, 10]);
-
-        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 1, embedDim: input.shape.at(-1)! });
+        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 1, embedDim: input.shape.at(-1) });
         expect(Object.keys(attention.getConfig())).not.toBe(0);
-    })
-
-
+    });
     test("causal attention hard coded values", () => {
         // input and output shapes: [2, 3, 10]
         const input = tf.tensor([
             [[0.2109915, 0.6158954, 0.6012088, 0.9867562, 0.8728716, 0.7496274, 0.8173883, 0.2958342, 0.9650571, 0.2075207],
-            [0.2946285, 0.9779906, 0.3203818, 0.4037617, 0.3762881, 0.9863171, 0.6655593, 0.7707329, 0.3216831, 0.7984023],
-            [0.9080769, 0.0026282, 0.379492, 0.0162054, 0.1939302, 0.2201049, 0.8190675, 0.0203963, 0.0114392, 0.5015539]],
-
+                [0.2946285, 0.9779906, 0.3203818, 0.4037617, 0.3762881, 0.9863171, 0.6655593, 0.7707329, 0.3216831, 0.7984023],
+                [0.9080769, 0.0026282, 0.379492, 0.0162054, 0.1939302, 0.2201049, 0.8190675, 0.0203963, 0.0114392, 0.5015539]],
             [[0.6241482, 0.7631097, 0.6687831, 0.7259795, 0.0457698, 0.6889264, 0.0853676, 0.8697655, 0.3637198, 0.2105307],
-            [0.5221761, 0.4476321, 0.1244729, 0.8863543, 0.7319002, 0.2954829, 0.3200496, 0.0905503, 0.607977, 0.1309131],
-            [0.4693873, 0.4609751, 0.9170766, 0.7065565, 0.4795104, 0.3225758, 0.1353116, 0.7083887, 0.1928891, 0.967386]]
+                [0.5221761, 0.4476321, 0.1244729, 0.8863543, 0.7319002, 0.2954829, 0.3200496, 0.0905503, 0.607977, 0.1309131],
+                [0.4693873, 0.4609751, 0.9170766, 0.7065565, 0.4795104, 0.3225758, 0.1353116, 0.7083887, 0.1928891, 0.967386]]
         ]);
-
         const expected = tf.tensor([
             [[0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344],
-            [0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376],
-            [0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539]],
-
+                [0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376],
+                [0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539]],
             [[0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718],
-            [0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268],
-            [0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877]]
+                [0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268],
+                [0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877]]
         ]);
-
-
-        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 1, embedDim: input.shape.at(-1)!, causal: true });
+        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 1, embedDim: input.shape.at(-1), causal: true });
         attention.build(input.shape);
         attention.setWeights(attention.getWeights().map(weight => tf.onesLike(weight).mul(0.05)));
-
-        expect(expected.sub(attention.apply(input) as tf.Tensor).sum().dataSync()[0]).toBeLessThan(1e-6);
-    })
-
-
+        expect(expected.sub(attention.apply(input)).sum().dataSync()[0]).toBeLessThan(1e-6);
+    });
     test("non-causal attention hard coded values", () => {
         // input and output shapes: [2, 3, 10]
         const input = tf.tensor([
             [[0.2109915, 0.6158954, 0.6012088, 0.9867562, 0.8728716, 0.7496274, 0.8173883, 0.2958342, 0.9650571, 0.2075207],
-            [0.2946285, 0.9779906, 0.3203818, 0.4037617, 0.3762881, 0.9863171, 0.6655593, 0.7707329, 0.3216831, 0.7984023],
-            [0.9080769, 0.0026282, 0.379492, 0.0162054, 0.1939302, 0.2201049, 0.8190675, 0.0203963, 0.0114392, 0.5015539]],
-
+                [0.2946285, 0.9779906, 0.3203818, 0.4037617, 0.3762881, 0.9863171, 0.6655593, 0.7707329, 0.3216831, 0.7984023],
+                [0.9080769, 0.0026282, 0.379492, 0.0162054, 0.1939302, 0.2201049, 0.8190675, 0.0203963, 0.0114392, 0.5015539]],
             [[0.6241482, 0.7631097, 0.6687831, 0.7259795, 0.0457698, 0.6889264, 0.0853676, 0.8697655, 0.3637198, 0.2105307],
-            [0.5221761, 0.4476321, 0.1244729, 0.8863543, 0.7319002, 0.2954829, 0.3200496, 0.0905503, 0.607977, 0.1309131],
-            [0.4693873, 0.4609751, 0.9170766, 0.7065565, 0.4795104, 0.3225758, 0.1353116, 0.7083887, 0.1928891, 0.967386]]
+                [0.5221761, 0.4476321, 0.1244729, 0.8863543, 0.7319002, 0.2954829, 0.3200496, 0.0905503, 0.607977, 0.1309131],
+                [0.4693873, 0.4609751, 0.9170766, 0.7065565, 0.4795104, 0.3225758, 0.1353116, 0.7083887, 0.1928891, 0.967386]]
         ]);
-
-
         const expected = tf.tensor([
             [[0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344, 0.2055344],
-            [0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376],
-            [0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539]],
-
+                [0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376, 0.205376],
+                [0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539, 0.2042539]],
             [[0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718, 0.1966718],
-            [0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268],
-            [0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877]]
+                [0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268, 0.1966268],
+                [0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877, 0.1966877]]
         ]);
-
-        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 1, embedDim: input.shape.at(-1)!, causal: false });
+        const attention = new CachedRoPEMultiHeadAttention({ numHeads: 1, embedDim: input.shape.at(-1), causal: false });
         attention.build(input.shape);
         attention.setWeights(attention.getWeights().map(weight => tf.onesLike(weight).mul(0.05)));
-
-        expect(expected.sub(attention.apply(input) as tf.Tensor).sum().dataSync()[0]).toBeLessThan(1e-6);
+        expect(expected.sub(attention.apply(input)).sum().dataSync()[0]).toBeLessThan(1e-6);
     });
 });
+//# sourceMappingURL=multihead_attention.test.js.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"multihead_attention.test.js","sourceRoot":"","sources":["../../src/layers/multihead_attention.test.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAEvC,OAAO,EAAE,4BAA4B,EAAE,MAAM,0CAA0C,CAAC;AACxF,OAAO,EAAE,MAAM,IAAI,kBAAkB,EAAE,MAAM,SAAS,CAAC;AACvD,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,uGAAuG,EAAE,GAAG,EAAE;QAC7G,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACxF,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,QAAQ,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;IACjG,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,2BAA2B,EAAE,GAAG,EAAE;QACnC,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAE3C,MAAM,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACnG,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;QAErD,MAAM,MAAM,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,MAAM,EAAE,IAAI,EAAE,CAAC,CAAC;QAC9G,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,IAAI,CAAC,sFAAsF,EAAE,GAAG,EAAE;QAC9F,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAC3C,MAAM,GAAG,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QACzC,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAC3C,MAAM,oBAAoB,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,GAAG,EAAE,EAAE,CAAC,CAAC,CAAC;QAE5D,MAAM,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACnG,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,GAAG,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACjE,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,GAAG,EAAE,oBAAoB,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IAChF,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,mCAAmC,EAAE,GAAG,EAAE;QACzC,MAAM,UAAU,GAAG,EAAE,CAAC;QAEtB,MAAM,SAAS,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,UAAU,CAAC,CAAC,CAAC;QACpD,MAAM,IAAI,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,UAAU,CAAC,CAAC,CAAC;QAClD,MAAM,SAAS,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,UAAU,CAAC,CAAC,CAAC;QAE3D,MAAM,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,UAAU,EAAE,CAAC,CAAC;QAE1F,MAAM;QACN,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,SAAS,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACnD,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,SAAS,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACrD,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAE3E,KAAK;QACL,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,IAAI,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAClD,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACpD,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,IAAI,EAAE,IAAI,EAAE,IAAI,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAEhE,MAAM;QACN,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,SAAS,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACnD,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,SAAS,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACrD,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAE3E,MAAM;QACN,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,IAAI,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACtE,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,SAAS,EAAE,IAAI,EAAE,SAAS,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACtE,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,IAAI,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACtE,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,SAAS,EAAE,IAAI,EAAE,IAAI,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACjE,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,IAAI,EAAE,IAAI,EAAE,SAAS,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IACrE,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,IAAI,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QAEjG,KAAK;QACL,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,KAAK,EAAE,EAAE,WAAW,EAAE,SAAS,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAC/E,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACrD,gCAAgC;QAChC,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QAC7F,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAEnE,MAAM;QACN,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAC5C,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACxD,gCAAgC;QAChC,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QAC7F,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IAC1E,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,mBAAmB,EAAE,GAAG,EAAE;QAC3B,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAC3C,MAAM,GAAG,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QACzC,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAE3C,MAAM,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,MAAM,EAAE,IAAI,EAAE,CAAC,CAAC;QAEjH,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,IAAI,CAAC,KAAK,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAEtD,kBAAkB;QAClB,MAAM,CAAC,GAAG,EAAE,CAAC,SAAS,CAAC,IAAI,CAAC,CAAC,KAAK,EAAE,GAAG,EAAE,KAAK,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QAGpE,MAAM,MAAM,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAC5C,MAAM,IAAI,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAC1C,MAAM,MAAM,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,CAAC;QAE5C,MAAM,aAAa,GAAG,EAAE,CAAC,MAAM,CAAC,CAAC;gBAC7B,8EAA8E;gBAC9E,iFAAiF;gBACjF,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC;gBAC5B,CAAC,IAAI,EAAE,IAAI,EAAE,IAAI,EAAE,KAAK,EAAE;gBAC1B,CAAC,IAAI,EAAE,IAAI,EAAE,IAAI,EAAE,KAAK,EAAE;gBAC1B,CAAC,KAAK,EAAE,KAAK,EAAE,KAAK,EAAE,KAAK,CAAC;gBAC5B,CAAC,IAAI,EAAE,IAAI,EAAE,IAAI,EAAE,KAAK,EAAE;aAC7B,CAAC,CAAC,CAAC;QAEJ,MAAM,YAAY,GAAG,EAAE,CAAC,MAAM,CAAC;YAC3B,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,CAAC;YACrB,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,CAAC;YACrB,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,CAAC;YACrB,CAAC,CAAC,GAAG,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,EAAE,CAAC,EAAE,CAAC,CAAC;YACxB,CAAC,CAAC,GAAG,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,EAAE,CAAC,EAAE,CAAC,CAAC;SAC3B,CAAC,CAAA;QAEF,2DAA2D;QAC3D,MAAM,CAAC,GAAG,EAAE,CAAC,kBAAkB,CAAC,4BAA4B,CAAC,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,aAAa,CAAC,MAAM,CAAC,SAAS,CAAC,EAAE,IAAI,EAAE,IAAI,EAAE,GAAG,EAAE,IAAI,EAAE,EAAE,cAAc,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAC9K,yDAAyD;QACzD,MAAM,CAAC,GAAG,EAAE,CAAC,kBAAkB,CAAC,4BAA4B,CAAC,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,aAAa,CAAC,MAAM,CAAC,SAAS,CAAC,EAAE,IAAI,EAAE,kBAAkB,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,CAAE,EAAE,IAAI,CAAC,KAAK,CAAC,CAAC,CAAE,CAAC,EAAE,GAAG,EAAE,IAAI,EAAE,EAAE,cAAc,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAC9N,8EAA8E;QAC9E,MAAM,CAAC,GAAG,EAAE,CAAC,kBAAkB,CAAC,4BAA4B,CAAC,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,aAAa,CAAC,MAAM,CAAC,SAAS,CAAC,EAAE,IAAI,EAAE,IAAI,EAAE,CAAC,EAAE,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACzJ,kCAAkC;QAClC,MAAM,CAAC,GAAG,EAAE,CAAC,kBAAkB,CAAC,4BAA4B,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,IAAI,EAAE,YAAY,EAAE,IAAI,EAAE,GAAG,EAAE,IAAI,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;IAC7I,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,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC;QACnG,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,IAAI,CAAC,oCAAoC,EAAE,GAAG,EAAE;QAC5C,sCAAsC;QACtC,MAAM,KAAK,GAAG,EAAE,CAAC,MAAM,CAAC;YACpB,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC/G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC9G,CAAC,SAAS,EAAE,SAAS,EAAE,QAAQ,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC;YAE9G,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC/G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,QAAQ,EAAE,SAAS,CAAC;gBAC7G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,QAAQ,CAAC,CAAC;SACjH,CAAC,CAAC;QAEH,MAAM,QAAQ,GAAG,EAAE,CAAC,MAAM,CAAC;YACvB,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC/G,CAAC,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,CAAC;gBACpG,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC;YAE/G,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC/G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC9G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC;SAClH,CAAC,CAAC;QAGH,MAAM,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,MAAM,EAAE,IAAI,EAAE,CAAC,CAAC;QACjH,SAAS,CAAC,KAAK,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC;QAC7B,SAAS,CAAC,UAAU,CAAC,SAAS,CAAC,UAAU,EAAE,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,CAAC,EAAE,CAAC,QAAQ,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAE1F,MAAM,CAAC,QAAQ,CAAC,GAAG,CAAC,SAAS,CAAC,KAAK,CAAC,KAAK,CAAc,CAAC,CAAC,GAAG,EAAE,CAAC,QAAQ,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,YAAY,CAAC,IAAI,CAAC,CAAC;IACrG,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,wCAAwC,EAAE,GAAG,EAAE;QAChD,sCAAsC;QACtC,MAAM,KAAK,GAAG,EAAE,CAAC,MAAM,CAAC;YACpB,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC/G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC9G,CAAC,SAAS,EAAE,SAAS,EAAE,QAAQ,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC;YAE9G,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC/G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,QAAQ,EAAE,SAAS,CAAC;gBAC7G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,QAAQ,CAAC,CAAC;SACjH,CAAC,CAAC;QAGH,MAAM,QAAQ,GAAG,EAAE,CAAC,MAAM,CAAC;YACvB,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC/G,CAAC,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,EAAE,QAAQ,CAAC;gBACpG,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC;YAE/G,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC/G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC9G,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC;SAClH,CAAC,CAAC;QAEH,MAAM,SAAS,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,QAAQ,EAAE,KAAK,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,EAAE,MAAM,EAAE,KAAK,EAAE,CAAC,CAAC;QAClH,SAAS,CAAC,KAAK,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC;QAC7B,SAAS,CAAC,UAAU,CAAC,SAAS,CAAC,UAAU,EAAE,CAAC,GAAG,CAAC,MAAM,CAAC,EAAE,CAAC,EAAE,CAAC,QAAQ,CAAC,MAAM,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAE1F,MAAM,CAAC,QAAQ,CAAC,GAAG,CAAC,SAAS,CAAC,KAAK,CAAC,KAAK,CAAc,CAAC,CAAC,GAAG,EAAE,CAAC,QAAQ,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,YAAY,CAAC,IAAI,CAAC,CAAC;IACrG,CAAC,CAAC,CAAC;AACP,CAAC,CAAC,CAAC"}

package/dist/layers/positional_encoding.d.ts

@@ -0,0 +1,37 @@
+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";
+export interface PositionalEncodingArgs extends LayerArgs {
+    embedDim: number;
+    maxSequenceLength?: number;
+}
+/**
+ * This class implements the position encoding logic described in the
+ * 2017 paper "Attention Is All You Need".
+ *
+ * This layer is untrainable and accepts inputs of shape `[ batch, sequences, embedding dims ]`
+ * and adds positional encoding to return an output tensor of the same shape.
+ *
+ * @param embedDim the size of each token/word's embedding
+ * @param maxSequenceLength the max number of tokens (words) per input (sentence), default `5120`
+ */
+export declare class PositionalEncoding extends tf.layers.Layer {
+    static className: string;
+    private readonly maxSequenceLength;
+    private readonly embedDim;
+    private positionalEncodings;
+    constructor(args: PositionalEncodingArgs);
+    /**
+     * Forward propagation. Injects positional encoding to the input embeddings
+     */
+    call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[];
+    /**
+     * Generate the positional encoding from the paper Attention Is All You Need.
+     * Note that because the inner term of the position formula is the same for both even
+     * and odd indices, we only create half of it and apply sine and cosine individually.
+     */
+    build(inputShape: tf.Shape | tf.Shape[]): void;
+    computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[];
+    getConfig(): tf.serialization.ConfigDict;
+}
+//# sourceMappingURL=positional_encoding.d.ts.map

package/dist/layers/positional_encoding.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"positional_encoding.d.ts","sourceRoot":"","sources":["../../src/layers/positional_encoding.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,KAAK,SAAS,EAAE,MAAM,8CAA8C,CAAC;AAC9E,OAAO,EAAE,KAAK,MAAM,EAAE,MAAM,oCAAoC,CAAC;AAGjE,MAAM,WAAW,sBAAuB,SAAQ,SAAS;IAErD,QAAQ,EAAE,MAAM,CAAC;IAEjB,iBAAiB,CAAC,EAAE,MAAM,CAAC;CAC9B;AAGD;;;;;;;;;GASG;AACH,qBAAa,kBAAmB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACnD,MAAM,CAAC,SAAS,SAAwB;IACxC,OAAO,CAAC,QAAQ,CAAC,iBAAiB,CAAS;IAC3C,OAAO,CAAC,QAAQ,CAAC,QAAQ,CAAS;IAClC,OAAO,CAAC,mBAAmB,CAAmB;gBAGlC,IAAI,EAAE,sBAAsB;IAuBxC;;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;;;;OAIG;IACM,KAAK,CAAC,UAAU,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,IAAI;IAmD9C,kBAAkB,CAAC,UAAU,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE;IAK5E,SAAS,IAAI,EAAE,CAAC,aAAa,CAAC,UAAU;CAYpD"}

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

@@ -1,112 +1,81 @@
 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";
-
-
-export interface PositionalEncodingArgs extends LayerArgs {
-    // embedding size of each word/token, aka d_model from the paper
-    embedDim: number;
-    // the max length of each sentence, any more or less are truncated or padded
-    maxSequenceLength?: number;
-}
-
-
 /**
  * This class implements the position encoding logic described in the
  * 2017 paper "Attention Is All You Need".
- * 
+ *
  * This layer is untrainable and accepts inputs of shape `[ batch, sequences, embedding dims ]`
  * and adds positional encoding to return an output tensor of the same shape.
- * 
+ *
  * @param embedDim the size of each token/word's embedding
  * @param maxSequenceLength the max number of tokens (words) per input (sentence), default `5120`
  */
 export class PositionalEncoding extends tf.layers.Layer {
     static className = "PositionalEncoding";
-    private readonly maxSequenceLength: number;
-    private readonly embedDim: number;
-    private positionalEncodings: tf.LayerVariable;
-
-
-    constructor(args: PositionalEncodingArgs) {
+    maxSequenceLength;
+    embedDim;
+    positionalEncodings;
+    constructor(args) {
         super(args);
-
         this.maxSequenceLength = args.maxSequenceLength ?? 5120;
         this.embedDim = args.embedDim;
-
         if (this.maxSequenceLength < 1) {
             throw Error(`${this.getClassName()}::constructor ${this.name} maxSequenceLength` +
                 ` (${args.maxSequenceLength}) must be greater than 0`);
         }
-
         if (this.embedDim < 1) {
             throw Error(`${this.getClassName()}::constructor ${this.name} embedDim` +
                 ` (${args.embedDim}) must be greater than 0`);
         }
-
         // positional encodings are not trainable
-        this.positionalEncodings = this.addWeight('positional_encodings',
-            [this.maxSequenceLength, this.embedDim], "float32",
-            tf.initializers.zeros(), undefined, false);
+        this.positionalEncodings = this.addWeight('positional_encodings', [this.maxSequenceLength, this.embedDim], "float32", tf.initializers.zeros(), undefined, false);
     }
-
-
     /**
      * Forward propagation. Injects positional encoding to the input embeddings
      */
-    override call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[] {
+    call(inputs, kwargs) {
         // validate the input tensors
         const input = Array.isArray(inputs) ? inputs[0] : inputs;
-        const sequences = input.shape[1]!;
-
+        const sequences = input.shape[1];
         if (input.shape.length != 3 || input.shape[2] != this.embedDim) {
             throw Error(`${this.getClassName()}::call ${this.name} expected an input shape of` +
                 ` [batch, (up to ${this.maxSequenceLength}), ${this.embedDim}], instead got ${input.shape}`);
         }
-
         if (sequences > this.maxSequenceLength) {
             // unexpected sequence length
             throw Error(`${this.getClassName()}::call ${this.name} received an input with` +
                 ` sequence length (${sequences}) which is greater than the max sequence length` +
                 ` ${this.maxSequenceLength}`);
         }
-
         // perform forward propagation
         return tf.tidy(() => {
             return input.add(this.positionalEncodings.read()
                 .slice([0, 0], [sequences, this.embedDim]) // gets the first "sequences" rows
                 .expandDims(0)); // introduce the batch dimension and let add() broadcast it
-        })
+        });
     }
-
     /**
      * Generate the positional encoding from the paper Attention Is All You Need.
      * Note that because the inner term of the position formula is the same for both even
      * and odd indices, we only create half of it and apply sine and cosine individually.
      */
-    override build(inputShape: tf.Shape | tf.Shape[]): void {
+    build(inputShape) {
         tf.tidy(() => {
             const embedDimHalved = Math.ceil(this.embedDim / 2);
-
             // create the position matrix as [ 0, 1, 2, 3, etc ],
             // and broadcast it horizontally to match the number of embeddings,
             const numerator = tf.range(0, this.maxSequenceLength, 1)
                 .reshape([this.maxSequenceLength, 1])
                 // this creates an extra, unsued positional encoding column later on for odd embedding sizes
                 .broadcastTo([this.maxSequenceLength, embedDimHalved]);
-
             // the inner term's denominator's exponent's numerator is created as
             // [ 0, 0, 2, 2, 4, 4, etc ] ( technically [0, 2, 4] as explained above ) and not
             // [ 0, 2, 4, 6, 8, 10, etc ] because the even and odd indices are counted as pairs
             // when incrementing "i",
             // the denominator formula is 10_000^(2i/d_model) where each "i" is a sine cosine pair
             const denominator = tf.pow(10_000, tf.range(0, this.embedDim, 2).div(this.embedDim));
-
             const inner_term = numerator.div(denominator);
-
             const sine = tf.sin(inner_term);
             const cosine = tf.cos(inner_term);
-
             // horizontally interweave the sine and cosine columns together to form
             // [sin, cos, sin, cos, etc]
             // [sin, cos, sin, cos, etc]
@@ -115,44 +84,32 @@ export class PositionalEncoding extends tf.layers.Layer {
             const ALL_ROWS = -1;
             const ONE_COL = 1;
             const FIRST_ROW = 0;
-
             for (let targetCol = 0; targetCol < this.embedDim / 2; targetCol++) {
-                interweaved.push(sine.slice([FIRST_ROW, targetCol], [ALL_ROWS, ONE_COL]))
-
+                interweaved.push(sine.slice([FIRST_ROW, targetCol], [ALL_ROWS, ONE_COL]));
                 if (targetCol != Math.floor(this.embedDim / 2)) {
                     // for odd numbered embedDim sizes skip the last cosine column
                     // e.g. if embedDim = 5, create [ i=0 (sin), i=0 (cos), i=1 (sin), i=1 (cos), i=2 (sin) ]
                     // and the final i=2 (cos) is ignored
-                    interweaved.push(cosine.slice([FIRST_ROW, targetCol], [ALL_ROWS, ONE_COL]))
+                    interweaved.push(cosine.slice([FIRST_ROW, targetCol], [ALL_ROWS, ONE_COL]));
                 }
             }
-
             // add the positional encoding
             this.setWeights([tf.concat(interweaved, 1)]);
         });
-
         super.build(inputShape);
     }
-
-
-    override computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[] {
+    computeOutputShape(inputShape) {
         return inputShape;
     }
-
-
-    override getConfig(): tf.serialization.ConfigDict {
+    getConfig() {
         const base_config = super.getConfig();
-
         const config = {
             maxSequenceLength: this.maxSequenceLength,
             embedDim: this.embedDim,
-        }
-
+        };
         Object.assign(config, base_config);
-
         return config;
     }
 }
-
-
 tf.serialization.registerClass(PositionalEncoding);
+//# sourceMappingURL=positional_encoding.js.map

package/dist/layers/positional_encoding.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"positional_encoding.js","sourceRoot":"","sources":["../../src/layers/positional_encoding.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAavC;;;;;;;;;GASG;AACH,MAAM,OAAO,kBAAmB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACnD,MAAM,CAAC,SAAS,GAAG,oBAAoB,CAAC;IACvB,iBAAiB,CAAS;IAC1B,QAAQ,CAAS;IAC1B,mBAAmB,CAAmB;IAG9C,YAAY,IAA4B;QACpC,KAAK,CAAC,IAAI,CAAC,CAAC;QAEZ,IAAI,CAAC,iBAAiB,GAAG,IAAI,CAAC,iBAAiB,IAAI,IAAI,CAAC;QACxD,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,QAAQ,CAAC;QAE9B,IAAI,IAAI,CAAC,iBAAiB,GAAG,CAAC,EAAE,CAAC;YAC7B,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,iBAAiB,IAAI,CAAC,IAAI,oBAAoB;gBAC5E,KAAK,IAAI,CAAC,iBAAiB,0BAA0B,CAAC,CAAC;QAC/D,CAAC;QAED,IAAI,IAAI,CAAC,QAAQ,GAAG,CAAC,EAAE,CAAC;YACpB,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,iBAAiB,IAAI,CAAC,IAAI,WAAW;gBACnE,KAAK,IAAI,CAAC,QAAQ,0BAA0B,CAAC,CAAC;QACtD,CAAC;QAED,yCAAyC;QACzC,IAAI,CAAC,mBAAmB,GAAG,IAAI,CAAC,SAAS,CAAC,sBAAsB,EAC5D,CAAC,IAAI,CAAC,iBAAiB,EAAE,IAAI,CAAC,QAAQ,CAAC,EAAE,SAAS,EAClD,EAAE,CAAC,YAAY,CAAC,KAAK,EAAE,EAAE,SAAS,EAAE,KAAK,CAAC,CAAC;IACnD,CAAC;IAGD;;OAEG;IACM,IAAI,CAAC,MAA+B,EAAE,MAAc;QACzD,6BAA6B;QAC7B,MAAM,KAAK,GAAG,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC;QACzD,MAAM,SAAS,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAE,CAAC;QAElC,IAAI,KAAK,CAAC,KAAK,CAAC,MAAM,IAAI,CAAC,IAAI,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,IAAI,IAAI,CAAC,QAAQ,EAAE,CAAC;YAC7D,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,UAAU,IAAI,CAAC,IAAI,6BAA6B;gBAC9E,mBAAmB,IAAI,CAAC,iBAAiB,MAAM,IAAI,CAAC,QAAQ,kBAAkB,KAAK,CAAC,KAAK,EAAE,CAAC,CAAC;QACrG,CAAC;QAED,IAAI,SAAS,GAAG,IAAI,CAAC,iBAAiB,EAAE,CAAC;YACrC,6BAA6B;YAC7B,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,UAAU,IAAI,CAAC,IAAI,yBAAyB;gBAC1E,qBAAqB,SAAS,iDAAiD;gBAC/E,IAAI,IAAI,CAAC,iBAAiB,EAAE,CAAC,CAAC;QACtC,CAAC;QAED,8BAA8B;QAC9B,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,OAAO,KAAK,CAAC,GAAG,CAAC,IAAI,CAAC,mBAAmB,CAAC,IAAI,EAAE;iBAC3C,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,SAAS,EAAE,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC,kCAAkC;iBAC5E,UAAU,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,2DAA2D;QACpF,CAAC,CAAC,CAAA;IACN,CAAC;IAED;;;;OAIG;IACM,KAAK,CAAC,UAAiC;QAC5C,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YACT,MAAM,cAAc,GAAG,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,QAAQ,GAAG,CAAC,CAAC,CAAC;YAEpD,qDAAqD;YACrD,mEAAmE;YACnE,MAAM,SAAS,GAAG,EAAE,CAAC,KAAK,CAAC,CAAC,EAAE,IAAI,CAAC,iBAAiB,EAAE,CAAC,CAAC;iBACnD,OAAO,CAAC,CAAC,IAAI,CAAC,iBAAiB,EAAE,CAAC,CAAC,CAAC;gBACrC,4FAA4F;iBAC3F,WAAW,CAAC,CAAC,IAAI,CAAC,iBAAiB,EAAE,cAAc,CAAC,CAAC,CAAC;YAE3D,oEAAoE;YACpE,iFAAiF;YACjF,mFAAmF;YACnF,yBAAyB;YACzB,sFAAsF;YACtF,MAAM,WAAW,GAAG,EAAE,CAAC,GAAG,CAAC,MAAM,EAAE,EAAE,CAAC,KAAK,CAAC,CAAC,EAAE,IAAI,CAAC,QAAQ,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC;YAErF,MAAM,UAAU,GAAG,SAAS,CAAC,GAAG,CAAC,WAAW,CAAC,CAAC;YAE9C,MAAM,IAAI,GAAG,EAAE,CAAC,GAAG,CAAC,UAAU,CAAC,CAAC;YAChC,MAAM,MAAM,GAAG,EAAE,CAAC,GAAG,CAAC,UAAU,CAAC,CAAC;YAElC,uEAAuE;YACvE,4BAA4B;YAC5B,4BAA4B;YAC5B,MAAM;YACN,MAAM,WAAW,GAAG,EAAE,CAAC;YACvB,MAAM,QAAQ,GAAG,CAAC,CAAC,CAAC;YACpB,MAAM,OAAO,GAAG,CAAC,CAAC;YAClB,MAAM,SAAS,GAAG,CAAC,CAAC;YAEpB,KAAK,IAAI,SAAS,GAAG,CAAC,EAAE,SAAS,GAAG,IAAI,CAAC,QAAQ,GAAG,CAAC,EAAE,SAAS,EAAE,EAAE,CAAC;gBACjE,WAAW,CAAC,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC,SAAS,EAAE,SAAS,CAAC,EAAE,CAAC,QAAQ,EAAE,OAAO,CAAC,CAAC,CAAC,CAAA;gBAEzE,IAAI,SAAS,IAAI,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,QAAQ,GAAG,CAAC,CAAC,EAAE,CAAC;oBAC7C,8DAA8D;oBAC9D,yFAAyF;oBACzF,qCAAqC;oBACrC,WAAW,CAAC,IAAI,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,SAAS,EAAE,SAAS,CAAC,EAAE,CAAC,QAAQ,EAAE,OAAO,CAAC,CAAC,CAAC,CAAA;gBAC/E,CAAC;YACL,CAAC;YAED,8BAA8B;YAC9B,IAAI,CAAC,UAAU,CAAC,CAAC,EAAE,CAAC,MAAM,CAAC,WAAW,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC;QACjD,CAAC,CAAC,CAAC;QAEH,KAAK,CAAC,KAAK,CAAC,UAAU,CAAC,CAAC;IAC5B,CAAC;IAGQ,kBAAkB,CAAC,UAAiC;QACzD,OAAO,UAAU,CAAC;IACtB,CAAC;IAGQ,SAAS;QACd,MAAM,WAAW,GAAG,KAAK,CAAC,SAAS,EAAE,CAAC;QAEtC,MAAM,MAAM,GAAG;YACX,iBAAiB,EAAE,IAAI,CAAC,iBAAiB;YACzC,QAAQ,EAAE,IAAI,CAAC,QAAQ;SAC1B,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/positional_encoding.test.d.ts

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

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

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