@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

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

@@ -1,113 +1,95 @@
 import * as tf from '@tensorflow/tfjs';
-
 import { PositionalEncoding } from '@/layers/positional_encoding';
-
 // 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", () => {
     it("should fail to instantiate a layer", () => {
         expect(() => new PositionalEncoding({ maxSequenceLength: 3, embedDim: 0 })).toThrow();
         expect(() => new PositionalEncoding({ maxSequenceLength: 3, embedDim: -1 })).toThrow();
         expect(() => new PositionalEncoding({ maxSequenceLength: 0, embedDim: 32 })).toThrow();
         expect(() => new PositionalEncoding({ maxSequenceLength: -1, embedDim: 32 })).toThrow();
-    })
-
-
+    });
     test("successfull forward calls", () => {
         const embed_dims = 32;
         const sequences = 4;
         const input = tf.randomUniform([2, sequences, embed_dims]);
-
         const positional = new PositionalEncoding({ embedDim: embed_dims });
         expect(() => positional.apply(input)).not.toThrow();
         expect(() => positional.apply([input])).not.toThrow();
         expect(positional.computeOutputShape(input.shape)).toEqual(input.shape);
-    })
-
-
+    });
     it("should throw when input sequences are too large, embedding dims don't match, input aren't rank 3", () => {
         const sequences_too_long = tf.randomUniform([100, 32]);
         const embeddings_too_large = tf.randomUniform([32, 100]);
         const wrong_rank = tf.randomUniform([10, 32, 32]);
-
         const positional = new PositionalEncoding({ maxSequenceLength: 10, embedDim: 32 });
-
         expect(() => positional.apply(sequences_too_long)).toThrow();
         expect(() => positional.apply(embeddings_too_large)).toThrow();
         expect(() => positional.apply(wrong_rank)).toThrow();
-    })
-
-
+    });
     it("should return a non-empty config dict", () => {
         const attention = new PositionalEncoding({ embedDim: 32 });
         expect(Object.keys(attention.getConfig())).not.toBe(0);
-    })
-
-
+    });
     // PyTorch implementation at found at
     // https://pytorch-tutorials-preview.netlify.app/beginner/transformer_tutorial.html
     it("should be within 1e-6 of PyTorch's implementation", () => {
         const pytorch_embed4 = tf.tensor([
             [[0.0000000, 1.0000000, 0.0000000, 1.0000000],
-            [0.8414710, 0.5403023, 0.0099998, 0.9999500],
-            [0.9092974, -0.4161468, 0.0199987, 0.9998000],
-            [0.1411200, -0.9899925, 0.0299955, 0.9995500],
-            [-0.7568025, -0.6536436, 0.0399893, 0.9992001],
-            [-0.9589243, 0.2836622, 0.0499792, 0.9987503],
-            [-0.2794155, 0.9601703, 0.0599640, 0.9982005],
-            [0.6569866, 0.7539023, 0.0699428, 0.9975510],
-            [0.9893582, -0.1455000, 0.0799147, 0.9968017],
-            [0.4121185, -0.9111302, 0.0898785, 0.9959527]]]);
-
+                [0.8414710, 0.5403023, 0.0099998, 0.9999500],
+                [0.9092974, -0.4161468, 0.0199987, 0.9998000],
+                [0.1411200, -0.9899925, 0.0299955, 0.9995500],
+                [-0.7568025, -0.6536436, 0.0399893, 0.9992001],
+                [-0.9589243, 0.2836622, 0.0499792, 0.9987503],
+                [-0.2794155, 0.9601703, 0.0599640, 0.9982005],
+                [0.6569866, 0.7539023, 0.0699428, 0.9975510],
+                [0.9893582, -0.1455000, 0.0799147, 0.9968017],
+                [0.4121185, -0.9111302, 0.0898785, 0.9959527]]
+        ]);
         const pytorch_embed8 = tf.tensor([
             [[0.0000000e+00, 1.0000000e+00, 0.0000000e+00, 1.0000000e+00,
-                0.0000000e+00, 1.0000000e+00, 0.0000000e+00, 1.0000000e+00],
-            [8.4147096e-01, 5.4030234e-01, 9.9833414e-02, 9.9500418e-01,
-                9.9998331e-03, 9.9994999e-01, 9.9999981e-04, 9.9999952e-01],
-            [9.0929741e-01, -4.1614684e-01, 1.9866931e-01, 9.8006660e-01,
-                1.9998666e-02, 9.9980003e-01, 1.9999985e-03, 9.9999803e-01],
-            [1.4112000e-01, -9.8999250e-01, 2.9552019e-01, 9.5533651e-01,
-                2.9995499e-02, 9.9955004e-01, 2.9999954e-03, 9.9999553e-01],
-            [-7.5680250e-01, -6.5364361e-01, 3.8941833e-01, 9.2106098e-01,
-                3.9989334e-02, 9.9920011e-01, 3.9999890e-03, 9.9999201e-01],
-            [-9.5892429e-01, 2.8366220e-01, 4.7942552e-01, 8.7758255e-01,
-                4.9979165e-02, 9.9875027e-01, 4.9999789e-03, 9.9998754e-01],
-            [-2.7941549e-01, 9.6017027e-01, 5.6464243e-01, 8.2533562e-01,
-                5.9964005e-02, 9.9820054e-01, 5.9999637e-03, 9.9998200e-01],
-            [6.5698659e-01, 7.5390226e-01, 6.4421761e-01, 7.6484221e-01,
-                6.9942847e-02, 9.9755102e-01, 6.9999420e-03, 9.9997550e-01],
-            [9.8935825e-01, -1.4550003e-01, 7.1735609e-01, 6.9670677e-01,
-                7.9914689e-02, 9.9680167e-01, 7.9999138e-03, 9.9996799e-01],
-            [4.1211849e-01, -9.1113025e-01, 7.8332686e-01, 6.2160999e-01,
-                8.9878544e-02, 9.9595273e-01, 8.9998785e-03, 9.9995953e-01]]]);
-
+                    0.0000000e+00, 1.0000000e+00, 0.0000000e+00, 1.0000000e+00],
+                [8.4147096e-01, 5.4030234e-01, 9.9833414e-02, 9.9500418e-01,
+                    9.9998331e-03, 9.9994999e-01, 9.9999981e-04, 9.9999952e-01],
+                [9.0929741e-01, -4.1614684e-01, 1.9866931e-01, 9.8006660e-01,
+                    1.9998666e-02, 9.9980003e-01, 1.9999985e-03, 9.9999803e-01],
+                [1.4112000e-01, -9.8999250e-01, 2.9552019e-01, 9.5533651e-01,
+                    2.9995499e-02, 9.9955004e-01, 2.9999954e-03, 9.9999553e-01],
+                [-7.5680250e-01, -6.5364361e-01, 3.8941833e-01, 9.2106098e-01,
+                    3.9989334e-02, 9.9920011e-01, 3.9999890e-03, 9.9999201e-01],
+                [-9.5892429e-01, 2.8366220e-01, 4.7942552e-01, 8.7758255e-01,
+                    4.9979165e-02, 9.9875027e-01, 4.9999789e-03, 9.9998754e-01],
+                [-2.7941549e-01, 9.6017027e-01, 5.6464243e-01, 8.2533562e-01,
+                    5.9964005e-02, 9.9820054e-01, 5.9999637e-03, 9.9998200e-01],
+                [6.5698659e-01, 7.5390226e-01, 6.4421761e-01, 7.6484221e-01,
+                    6.9942847e-02, 9.9755102e-01, 6.9999420e-03, 9.9997550e-01],
+                [9.8935825e-01, -1.4550003e-01, 7.1735609e-01, 6.9670677e-01,
+                    7.9914689e-02, 9.9680167e-01, 7.9999138e-03, 9.9996799e-01],
+                [4.1211849e-01, -9.1113025e-01, 7.8332686e-01, 6.2160999e-01,
+                    8.9878544e-02, 9.9595273e-01, 8.9998785e-03, 9.9995953e-01]]
+        ]);
         const positional4 = new PositionalEncoding({ embedDim: 4, maxSequenceLength: 10 });
         positional4.build([]);
-
         const positional8 = new PositionalEncoding({ embedDim: 8, maxSequenceLength: 10 });
         positional8.build([]);
-
         const margin_of_error = 1e-6;
-
         // the difference between this and PyTorch's implementation
         //should be insignificantly small
-        expect((positional4.getWeights()[0]
+        expect(positional4.getWeights()[0]
             .sub(pytorch_embed4)
             .abs()
-            .arraySync() as [])
+            .arraySync()
             .flat(2)
             .filter(i => i > margin_of_error)
             .length).toBe(0);
-
-        expect((positional8.getWeights()[0]
+        expect(positional8.getWeights()[0]
             .sub(pytorch_embed8)
             .abs()
-            .arraySync() as [])
+            .arraySync()
             .flat(2)
             .filter(i => i > margin_of_error)
             .length).toBe(0);
     });
 });
+//# sourceMappingURL=positional_encoding.test.js.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"positional_encoding.test.js","sourceRoot":"","sources":["../../src/layers/positional_encoding.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,oCAAoC,EAAE,GAAG,EAAE;QAC1C,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,kBAAkB,CAAC,EAAE,iBAAiB,EAAE,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACtF,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,kBAAkB,CAAC,EAAE,iBAAiB,EAAE,CAAC,EAAE,QAAQ,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACvF,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,kBAAkB,CAAC,EAAE,iBAAiB,EAAE,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACvF,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,kBAAkB,CAAC,EAAE,iBAAiB,EAAE,CAAC,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;IAC5F,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,2BAA2B,EAAE,GAAG,EAAE;QACnC,MAAM,UAAU,GAAG,EAAE,CAAC;QACtB,MAAM,SAAS,GAAG,CAAC,CAAC;QACpB,MAAM,KAAK,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,SAAS,EAAE,UAAU,CAAC,CAAC,CAAC;QAE3D,MAAM,UAAU,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,UAAU,EAAE,CAAC,CAAC;QACpE,MAAM,CAAC,GAAG,EAAE,CAAC,UAAU,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACpD,MAAM,CAAC,GAAG,EAAE,CAAC,UAAU,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;QACtD,MAAM,CAAC,UAAU,CAAC,kBAAkB,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,OAAO,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC;IAC5E,CAAC,CAAC,CAAA;IAGF,EAAE,CAAC,kGAAkG,EAAE,GAAG,EAAE;QACxG,MAAM,kBAAkB,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,GAAG,EAAE,EAAE,CAAC,CAAC,CAAC;QACvD,MAAM,oBAAoB,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,EAAE,EAAE,GAAG,CAAC,CAAC,CAAC;QACzD,MAAM,UAAU,GAAG,EAAE,CAAC,aAAa,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC;QAElD,MAAM,UAAU,GAAG,IAAI,kBAAkB,CAAC,EAAE,iBAAiB,EAAE,EAAE,EAAE,QAAQ,EAAE,EAAE,EAAE,CAAC,CAAC;QAEnF,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,oBAAoB,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAC/D,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,kBAAkB,CAAC,EAAE,QAAQ,EAAE,EAAE,EAAE,CAAC,CAAC;QAC3D,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,qCAAqC;IACrC,mFAAmF;IACnF,EAAE,CAAC,mDAAmD,EAAE,GAAG,EAAE;QACzD,MAAM,cAAc,GAAG,EAAE,CAAC,MAAM,CAAC;YAC7B,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC7C,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC5C,CAAC,SAAS,EAAE,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC7C,CAAC,SAAS,EAAE,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC7C,CAAC,CAAC,SAAS,EAAE,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC9C,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC7C,CAAC,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC7C,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC5C,CAAC,SAAS,EAAE,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;gBAC7C,CAAC,SAAS,EAAE,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC,CAAC;SAAC,CAAC,CAAC;QAErD,MAAM,cAAc,GAAG,EAAE,CAAC,MAAM,CAAC;YAC7B,CAAC,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa;oBACxD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa;oBACvD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,aAAa,EAAE,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa;oBACxD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,aAAa,EAAE,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa;oBACxD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,CAAC,aAAa,EAAE,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa;oBACzD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa;oBACxD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa;oBACxD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa;oBACvD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,aAAa,EAAE,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa;oBACxD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC;gBAC/D,CAAC,aAAa,EAAE,CAAC,aAAa,EAAE,aAAa,EAAE,aAAa;oBACxD,aAAa,EAAE,aAAa,EAAE,aAAa,EAAE,aAAa,CAAC,CAAC;SAAC,CAAC,CAAC;QAEvE,MAAM,WAAW,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,iBAAiB,EAAE,EAAE,EAAE,CAAC,CAAC;QACnF,WAAW,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;QAEtB,MAAM,WAAW,GAAG,IAAI,kBAAkB,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,iBAAiB,EAAE,EAAE,EAAE,CAAC,CAAC;QACnF,WAAW,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;QAEtB,MAAM,eAAe,GAAG,IAAI,CAAC;QAE7B,2DAA2D;QAC3D,iCAAiC;QACjC,MAAM,CAAE,WAAW,CAAC,UAAU,EAAE,CAAC,CAAC,CAAC;aAC9B,GAAG,CAAC,cAAc,CAAC;aACnB,GAAG,EAAE;aACL,SAAS,EAAS;aAClB,IAAI,CAAC,CAAC,CAAC;aACP,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,GAAG,eAAe,CAAC;aAChC,MAAM,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;QAErB,MAAM,CAAE,WAAW,CAAC,UAAU,EAAE,CAAC,CAAC,CAAC;aAC9B,GAAG,CAAC,cAAc,CAAC;aACnB,GAAG,EAAE;aACL,SAAS,EAAS;aAClB,IAAI,CAAC,CAAC,CAAC;aACP,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,GAAG,eAAe,CAAC;aAChC,MAAM,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;IACzB,CAAC,CAAC,CAAC;AACP,CAAC,CAAC,CAAC"}

package/dist/layers/rotary_position_embedding.d.ts

@@ -0,0 +1,39 @@
+import * as tf from "@tensorflow/tfjs";
+import { type LayerArgs } from "@tensorflow/tfjs-layers/dist/engine/topology";
+export declare function applyRope(x: tf.Tensor, dim: number, cosine_cache: tf.Tensor, sine_cache: tf.Tensor): tf.Tensor<tf.Rank>;
+export declare function rotateHalf(x: tf.Tensor, dim: number): tf.Tensor;
+export declare function createRoPECache(dim: number, max_sequence_length: number, theta?: number): tf.Tensor<tf.Rank>[];
+export interface RotaryPositionEmbeddingArgs extends LayerArgs {
+    /**
+     * The dimension of each head (rounded down), e.g. `Math.floor(embedDim / numHeads)`
+     */
+    dim: number;
+    /**
+     * The RoPE cache will be pre-calculated up to the max sequence length, and re-caculated as needed. Defaults to `4096`.
+     */
+    maxSequenceLength?: number;
+    /**
+     * The base for the geometric progression used to compute the rotation angles. Defaults to `10_000`.
+     */
+    theta?: number;
+}
+/**
+ * Implements RoPE from the RoFormer: Enhanced Transformer with Rotary Position Embedding paper
+ * Inspired by: https://meta-pytorch.org/torchtune/stable/_modules/torchtune/modules/position_embeddings.html#RotaryPositionalEmbeddings
+ */
+export declare class RotaryPositionEmbedding extends tf.layers.Layer {
+    static className: string;
+    protected dim: number;
+    protected max_sequence_length: number;
+    protected theta: number;
+    protected cosine_cache: tf.LayerVariable;
+    protected sine_cache: tf.LayerVariable;
+    constructor({ dim, maxSequenceLength, theta, ...args }: RotaryPositionEmbeddingArgs);
+    call(inputs: tf.Tensor | tf.Tensor[], kwargs: any): tf.Tensor | tf.Tensor[];
+    build(input_shape: tf.Shape | tf.Shape[]): void;
+    /**
+     * Output shape: [batch, head, sequence, head_dim]
+     */
+    computeOutputShape(input_shape: tf.Shape | tf.Shape[]): tf.Shape;
+}
+//# sourceMappingURL=rotary_position_embedding.d.ts.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"rotary_position_embedding.d.ts","sourceRoot":"","sources":["../../src/layers/rotary_position_embedding.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,KAAK,SAAS,EAAE,MAAM,8CAA8C,CAAC;AAG9E,wBAAgB,SAAS,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,EAAE,GAAG,EAAE,MAAM,EAAE,YAAY,EAAE,EAAE,CAAC,MAAM,EAAE,UAAU,EAAE,EAAE,CAAC,MAAM,sBAalG;AAGD,wBAAgB,UAAU,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,EAAE,GAAG,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM,CAgB/D;AAGD,wBAAgB,eAAe,CAAC,GAAG,EAAE,MAAM,EAAE,mBAAmB,EAAE,MAAM,EAAE,KAAK,GAAE,MAAe,wBAqB/F;AAGD,MAAM,WAAW,2BAA4B,SAAQ,SAAS;IAC1D;;OAEG;IACH,GAAG,EAAE,MAAM,CAAC;IACZ;;OAEG;IACH,iBAAiB,CAAC,EAAE,MAAM,CAAC;IAC3B;;OAEG;IACH,KAAK,CAAC,EAAE,MAAM,CAAC;CAClB;AAGD;;;GAGG;AACH,qBAAa,uBAAwB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACxD,MAAM,CAAC,SAAS,SAA6B;IAE7C,SAAS,CAAC,GAAG,EAAE,MAAM,CAAC;IACtB,SAAS,CAAC,mBAAmB,EAAE,MAAM,CAAC;IACtC,SAAS,CAAC,KAAK,EAAE,MAAM,CAAC;IAGxB,SAAS,CAAC,YAAY,EAAE,EAAE,CAAC,aAAa,CAAC;IACzC,SAAS,CAAC,UAAU,EAAE,EAAE,CAAC,aAAa,CAAC;gBAE3B,EAAE,GAAG,EAAE,iBAAwB,EAAE,KAAc,EAAE,GAAG,IAAI,EAAE,EAAE,2BAA2B;IAqB1F,IAAI,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE,EAAE,MAAM,EAAE,GAAG,GAAG,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE;IAkB3E,KAAK,CAAC,WAAW,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE;IAmBjD;;OAEG;IACI,kBAAkB,CAAC,WAAW,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE;CAK/D"}

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

@@ -1,163 +1,99 @@
 import * as tf from "@tensorflow/tfjs";
-import { type LayerArgs } from "@tensorflow/tfjs-layers/dist/engine/topology";
-
-
-export function applyRope(x: tf.Tensor, dim: number, cosine_cache: tf.Tensor, sine_cache: tf.Tensor) {
+export function applyRope(x, dim, cosine_cache, sine_cache) {
     return tf.tidy(() => {
-        const seq_length = x.shape[2]!;
-
+        const seq_length = x.shape[2];
         // get a slice of the pre-computed cache, up to the input's sequence length
         const cosine = cosine_cache.slice([0, 0, 0, 0], [1, 1, seq_length, dim]);
         const sine = sine_cache.slice([0, 0, 0, 0], [1, 1, seq_length, dim]);
-
         // apply RoPE formula (x1 * cosine) + (rotate(-x2) * sine)
         const rotated_x = rotateHalf(x, dim);
-
         return tf.add(tf.mul(x, cosine), tf.mul(rotated_x, sine));
     });
 }
-
-
-export function rotateHalf(x: tf.Tensor, dim: number): tf.Tensor {
+export function rotateHalf(x, dim) {
     return tf.tidy(() => {
         // reshape the last dimension such that adjacent coordinates are paired together
         // [x1, x2, x3, x4] -> [[x1, x2], [x3, x4]]
         // the leading dimensions are flattened because TFJS has issues during
         // backpropagation with 5D slicing
         const reshaped = x.reshape([-1, dim / 2, 2]);
-
         const x1 = reshaped.slice([0, 0, 0], [-1, -1, 1]);
         const x2 = reshaped.slice([0, 0, 1], [-1, -1, 1]);
-
         // [x1, x2] -> [-x2, x1]
         const rotated = tf.concat([tf.neg(x2), x1], -1);
-
         return rotated.reshape(x.shape);
     });
 }
-
-
-export function createRoPECache(dim: number, max_sequence_length: number, theta: number = 10_000) {
+export function createRoPECache(dim, max_sequence_length, theta = 10_000) {
     return tf.tidy(() => {
         // [dim]
-        const inv_frequencies = tf.div<tf.Tensor1D>(1, tf.pow(
-            theta,
-            tf.range(0, Math.floor(dim / 2) * 2, 2, "float32").div(dim)));
-
+        const inv_frequencies = tf.div(1, tf.pow(theta, tf.range(0, Math.floor(dim / 2) * 2, 2, "float32").div(dim)));
         // [max_sequene_length]
         const sequence_indices = tf.range(0, max_sequence_length);
         // 
         const freq = tf.outerProduct(sequence_indices, inv_frequencies);
-
         // cache final shape [max_sequence_length, dim]
         const freq_pairs = tf.stack([freq, freq], -1)
             .reshape([max_sequence_length, dim]);
-
         return [
             tf.keep(tf.cos(freq_pairs).expandDims(0).expandDims(0)),
             tf.keep(tf.sin(freq_pairs).expandDims(0).expandDims(0))
-        ]
+        ];
     });
 }
-
-
-export interface RotaryPositionEmbeddingArgs extends LayerArgs {
-    /**
-     * The dimension of each head (rounded down), e.g. `Math.floor(embedDim / numHeads)`
-     */
-    dim: number,
-    /**
-     * The RoPE cache will be pre-calculated up to the max sequence length, and re-caculated as needed. Defaults to `4096`.
-     */
-    maxSequenceLength?: number,
-    /**
-     * The base for the geometric progression used to compute the rotation angles. Defaults to `10_000`.
-     */
-    theta?: number,
-}
-
-
 /**
  * Implements RoPE from the RoFormer: Enhanced Transformer with Rotary Position Embedding paper
  * Inspired by: https://meta-pytorch.org/torchtune/stable/_modules/torchtune/modules/position_embeddings.html#RotaryPositionalEmbeddings
  */
 export class RotaryPositionEmbedding extends tf.layers.Layer {
     static className = "RotaryPositionEmbedding";
-
-    protected dim: number;
-    protected max_sequence_length: number;
-    protected theta: number;
-
+    dim;
+    max_sequence_length;
+    theta;
     // cached sine and cosine frequencies, untrainable weights
-    protected cosine_cache: tf.LayerVariable;
-    protected sine_cache: tf.LayerVariable;
-
-    constructor({ dim, maxSequenceLength = 4096, theta = 10_000, ...args }: RotaryPositionEmbeddingArgs) {
+    cosine_cache;
+    sine_cache;
+    constructor({ dim, maxSequenceLength = 4096, theta = 10_000, ...args }) {
         super(args);
-
         if (dim % 2 !== 0) {
             throw Error(`${this.getClassName()}::constructor ${this.name} expected dim to be even, got ${dim}`);
         }
-
         this.dim = dim;
         this.max_sequence_length = maxSequenceLength;
         this.theta = theta;
-
-        this.cosine_cache = this.addWeight("sine_cache",
-            [1, 1, maxSequenceLength, Math.floor(this.dim)],
-            "float32", tf.initializers.zeros(), undefined, false);
-
-        this.sine_cache = this.addWeight("cosine_cache",
-            [1, 1, maxSequenceLength, Math.floor(this.dim)],
-            "float32", tf.initializers.zeros(), undefined, false);
+        this.cosine_cache = this.addWeight("sine_cache", [1, 1, maxSequenceLength, Math.floor(this.dim)], "float32", tf.initializers.zeros(), undefined, false);
+        this.sine_cache = this.addWeight("cosine_cache", [1, 1, maxSequenceLength, Math.floor(this.dim)], "float32", tf.initializers.zeros(), undefined, false);
     }
-
-
-    override call(inputs: tf.Tensor | tf.Tensor[], kwargs: any): tf.Tensor | tf.Tensor[] {
+    call(inputs, kwargs) {
         const shape = Array.isArray(inputs) ? inputs[0].shape : inputs.shape;
         const seq_length = shape[2];
-
         if (seq_length > this.max_sequence_length) {
             // expand cache to the nearest power of 2
             this.max_sequence_length = Math.pow(2, Math.ceil(Math.log2(seq_length)));
             this.build([]);
         }
-
-        return applyRope(
-            Array.isArray(inputs) ? inputs[0] : inputs,
-            this.dim,
-            this.cosine_cache.read(),
-            this.sine_cache.read())
+        return applyRope(Array.isArray(inputs) ? inputs[0] : inputs, this.dim, this.cosine_cache.read(), this.sine_cache.read());
     }
-
-
-    override build(input_shape: tf.Shape | tf.Shape[]) {
-        const [cosine, sine] = createRoPECache(
-            this.dim, this.max_sequence_length, this.theta);
-
+    build(input_shape) {
+        const [cosine, sine] = createRoPECache(this.dim, this.max_sequence_length, this.theta);
         this.cosine_cache.dispose();
         this.sine_cache.dispose();
-
         this.cosine_cache = new tf.LayerVariable(cosine);
         this.sine_cache = new tf.LayerVariable(sine);
-
         this.nonTrainableWeights = [
             new tf.LayerVariable(cosine),
             new tf.LayerVariable(sine)
         ];
-
         this.setWeights([cosine, sine]);
     }
-
-
     /**
      * Output shape: [batch, head, sequence, head_dim]
      */
-    public computeOutputShape(input_shape: tf.Shape | tf.Shape[]) {
+    computeOutputShape(input_shape) {
         return Array.isArray(input_shape[0])
-            ? input_shape[0] as tf.Shape
-            : input_shape as tf.Shape;
+            ? input_shape[0]
+            : input_shape;
     }
 }
-
 tf.serialization.registerClass(RotaryPositionEmbedding);
+//# sourceMappingURL=rotary_position_embedding.js.map

package/dist/layers/rotary_position_embedding.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"rotary_position_embedding.js","sourceRoot":"","sources":["../../src/layers/rotary_position_embedding.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAIvC,MAAM,UAAU,SAAS,CAAC,CAAY,EAAE,GAAW,EAAE,YAAuB,EAAE,UAAqB;IAC/F,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;QAChB,MAAM,UAAU,GAAG,CAAC,CAAC,KAAK,CAAC,CAAC,CAAE,CAAC;QAE/B,2EAA2E;QAC3E,MAAM,MAAM,GAAG,YAAY,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,UAAU,EAAE,GAAG,CAAC,CAAC,CAAC;QACzE,MAAM,IAAI,GAAG,UAAU,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,UAAU,EAAE,GAAG,CAAC,CAAC,CAAC;QAErE,0DAA0D;QAC1D,MAAM,SAAS,GAAG,UAAU,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;QAErC,OAAO,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,EAAE,MAAM,CAAC,EAAE,EAAE,CAAC,GAAG,CAAC,SAAS,EAAE,IAAI,CAAC,CAAC,CAAC;IAC9D,CAAC,CAAC,CAAC;AACP,CAAC;AAGD,MAAM,UAAU,UAAU,CAAC,CAAY,EAAE,GAAW;IAChD,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;QAChB,gFAAgF;QAChF,2CAA2C;QAC3C,sEAAsE;QACtE,kCAAkC;QAClC,MAAM,QAAQ,GAAG,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,EAAE,GAAG,GAAG,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAE7C,MAAM,EAAE,GAAG,QAAQ,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAClD,MAAM,EAAE,GAAG,QAAQ,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAElD,wBAAwB;QACxB,MAAM,OAAO,GAAG,EAAE,CAAC,MAAM,CAAC,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC;QAEhD,OAAO,OAAO,CAAC,OAAO,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC;IACpC,CAAC,CAAC,CAAC;AACP,CAAC;AAGD,MAAM,UAAU,eAAe,CAAC,GAAW,EAAE,mBAA2B,EAAE,QAAgB,MAAM;IAC5F,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;QAChB,QAAQ;QACR,MAAM,eAAe,GAAG,EAAE,CAAC,GAAG,CAAc,CAAC,EAAE,EAAE,CAAC,GAAG,CACjD,KAAK,EACL,EAAE,CAAC,KAAK,CAAC,CAAC,EAAE,IAAI,CAAC,KAAK,CAAC,GAAG,GAAG,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,SAAS,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC;QAElE,uBAAuB;QACvB,MAAM,gBAAgB,GAAG,EAAE,CAAC,KAAK,CAAC,CAAC,EAAE,mBAAmB,CAAC,CAAC;QAC1D,GAAG;QACH,MAAM,IAAI,GAAG,EAAE,CAAC,YAAY,CAAC,gBAAgB,EAAE,eAAe,CAAC,CAAC;QAEhE,+CAA+C;QAC/C,MAAM,UAAU,GAAG,EAAE,CAAC,KAAK,CAAC,CAAC,IAAI,EAAE,IAAI,CAAC,EAAE,CAAC,CAAC,CAAC;aACxC,OAAO,CAAC,CAAC,mBAAmB,EAAE,GAAG,CAAC,CAAC,CAAC;QAEzC,OAAO;YACH,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,GAAG,CAAC,UAAU,CAAC,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;YACvD,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,GAAG,CAAC,UAAU,CAAC,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;SAC1D,CAAA;IACL,CAAC,CAAC,CAAC;AACP,CAAC;AAmBD;;;GAGG;AACH,MAAM,OAAO,uBAAwB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACxD,MAAM,CAAC,SAAS,GAAG,yBAAyB,CAAC;IAEnC,GAAG,CAAS;IACZ,mBAAmB,CAAS;IAC5B,KAAK,CAAS;IAExB,0DAA0D;IAChD,YAAY,CAAmB;IAC/B,UAAU,CAAmB;IAEvC,YAAY,EAAE,GAAG,EAAE,iBAAiB,GAAG,IAAI,EAAE,KAAK,GAAG,MAAM,EAAE,GAAG,IAAI,EAA+B;QAC/F,KAAK,CAAC,IAAI,CAAC,CAAC;QAEZ,IAAI,GAAG,GAAG,CAAC,KAAK,CAAC,EAAE,CAAC;YAChB,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,iBAAiB,IAAI,CAAC,IAAI,iCAAiC,GAAG,EAAE,CAAC,CAAC;QACxG,CAAC;QAED,IAAI,CAAC,GAAG,GAAG,GAAG,CAAC;QACf,IAAI,CAAC,mBAAmB,GAAG,iBAAiB,CAAC;QAC7C,IAAI,CAAC,KAAK,GAAG,KAAK,CAAC;QAEnB,IAAI,CAAC,YAAY,GAAG,IAAI,CAAC,SAAS,CAAC,YAAY,EAC3C,CAAC,CAAC,EAAE,CAAC,EAAE,iBAAiB,EAAE,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC,EAC/C,SAAS,EAAE,EAAE,CAAC,YAAY,CAAC,KAAK,EAAE,EAAE,SAAS,EAAE,KAAK,CAAC,CAAC;QAE1D,IAAI,CAAC,UAAU,GAAG,IAAI,CAAC,SAAS,CAAC,cAAc,EAC3C,CAAC,CAAC,EAAE,CAAC,EAAE,iBAAiB,EAAE,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,GAAG,CAAC,CAAC,EAC/C,SAAS,EAAE,EAAE,CAAC,YAAY,CAAC,KAAK,EAAE,EAAE,SAAS,EAAE,KAAK,CAAC,CAAC;IAC9D,CAAC;IAGQ,IAAI,CAAC,MAA+B,EAAE,MAAW;QACtD,MAAM,KAAK,GAAG,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,MAAM,CAAC,KAAK,CAAC;QACrE,MAAM,UAAU,GAAG,KAAK,CAAC,CAAC,CAAC,CAAC;QAE5B,IAAI,UAAU,GAAG,IAAI,CAAC,mBAAmB,EAAE,CAAC;YACxC,yCAAyC;YACzC,IAAI,CAAC,mBAAmB,GAAG,IAAI,CAAC,GAAG,CAAC,CAAC,EAAE,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,IAAI,CAAC,UAAU,CAAC,CAAC,CAAC,CAAC;YACzE,IAAI,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;QACnB,CAAC;QAED,OAAO,SAAS,CACZ,KAAK,CAAC,OAAO,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,MAAM,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,MAAM,EAC1C,IAAI,CAAC,GAAG,EACR,IAAI,CAAC,YAAY,CAAC,IAAI,EAAE,EACxB,IAAI,CAAC,UAAU,CAAC,IAAI,EAAE,CAAC,CAAA;IAC/B,CAAC;IAGQ,KAAK,CAAC,WAAkC;QAC7C,MAAM,CAAC,MAAM,EAAE,IAAI,CAAC,GAAG,eAAe,CAClC,IAAI,CAAC,GAAG,EAAE,IAAI,CAAC,mBAAmB,EAAE,IAAI,CAAC,KAAK,CAAC,CAAC;QAEpD,IAAI,CAAC,YAAY,CAAC,OAAO,EAAE,CAAC;QAC5B,IAAI,CAAC,UAAU,CAAC,OAAO,EAAE,CAAC;QAE1B,IAAI,CAAC,YAAY,GAAG,IAAI,EAAE,CAAC,aAAa,CAAC,MAAM,CAAC,CAAC;QACjD,IAAI,CAAC,UAAU,GAAG,IAAI,EAAE,CAAC,aAAa,CAAC,IAAI,CAAC,CAAC;QAE7C,IAAI,CAAC,mBAAmB,GAAG;YACvB,IAAI,EAAE,CAAC,aAAa,CAAC,MAAM,CAAC;YAC5B,IAAI,EAAE,CAAC,aAAa,CAAC,IAAI,CAAC;SAC7B,CAAC;QAEF,IAAI,CAAC,UAAU,CAAC,CAAC,MAAM,EAAE,IAAI,CAAC,CAAC,CAAC;IACpC,CAAC;IAGD;;OAEG;IACI,kBAAkB,CAAC,WAAkC;QACxD,OAAO,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,CAAC,CAAC,CAAC;YAChC,CAAC,CAAC,WAAW,CAAC,CAAC,CAAa;YAC5B,CAAC,CAAC,WAAuB,CAAC;IAClC,CAAC;;AAGL,EAAE,CAAC,aAAa,CAAC,aAAa,CAAC,uBAAuB,CAAC,CAAC"}

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

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

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

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

package/dist/layers/rotary_position_embedding.test.js

@@ -0,0 +1,88 @@
+import { RotaryPositionEmbedding } from "@/layers/rotary_position_embedding";
+import * as tf from "@tensorflow/tfjs";
+// disables warning for using the faster node backend,
+// https://github.com/tensorflow/tfjs/issues/5349#issuecomment-885170504
+tf.env().set('IS_NODE', false);
+describe("RotaryPositionEmbedding tests", () => {
+    test("create cache", async () => {
+        const rope = new RotaryPositionEmbedding({ dim: 8, maxSequenceLength: 15 });
+        rope.build([]);
+        const expected_cosine_cache = tf.tensor([[[
+                    [1, 1, 1, 1, 1, 1, 1, 1],
+                    [0.5403022766113281, 0.5403022766113281, 0.9950041770935059, 0.9950041770935059, 0.9999499917030334, 0.9999499917030334, 0.9999995231628418, 0.9999995231628418],
+                    [-0.416146844625473, -0.416146844625473, 0.9800665974617004, 0.9800665974617004, 0.9998000264167786, 0.9998000264167786, 0.9999979734420776, 0.9999979734420776],
+                    [-0.9899924993515015, -0.9899924993515015, 0.9553365111351013, 0.9553365111351013, 0.9995500445365906, 0.9995500445365906, 0.9999955296516418, 0.9999955296516418],
+                    [-0.6536436080932617, -0.6536436080932617, 0.9210609793663025, 0.9210609793663025, 0.9992001056671143, 0.9992001056671143, 0.9999920129776001, 0.9999920129776001],
+                    [0.28366219997406006, 0.28366219997406006, 0.8775825500488281, 0.8775825500488281, 0.9987502694129944, 0.9987502694129944, 0.9999874830245972, 0.9999874830245972],
+                    [0.9601702690124512, 0.9601702690124512, 0.8253356218338013, 0.8253356218338013, 0.998200535774231, 0.998200535774231, 0.9999819993972778, 0.9999819993972778],
+                    [0.7539022564888, 0.7539022564888, 0.7648422122001648, 0.7648422122001648, 0.9975510239601135, 0.9975510239601135, 0.9999755024909973, 0.9999755024909973],
+                    [-0.1455000340938568, -0.1455000340938568, 0.6967067122459412, 0.6967067122459412, 0.9968017339706421, 0.9968017339706421, 0.9999679923057556, 0.9999679923057556],
+                    [-0.9111302495002747, -0.9111302495002747, 0.6216099262237549, 0.6216099262237549, 0.9959527254104614, 0.9959527254104614, 0.9999595284461975, 0.9999595284461975],
+                    [-0.83907151222229, -0.83907151222229, 0.5403022766113281, 0.5403022766113281, 0.9950041770935059, 0.9950041770935059, 0.9999499917030334, 0.9999499917030334],
+                    [0.004425697959959507, 0.004425697959959507, 0.4535960853099823, 0.4535960853099823, 0.9939560890197754, 0.9939560890197754, 0.999939501285553, 0.999939501285553],
+                    [0.8438539505004883, 0.8438539505004883, 0.3623577058315277, 0.3623577058315277, 0.9928086400032043, 0.9928086400032043, 0.9999279975891113, 0.9999279975891113],
+                    [0.9074468016624451, 0.9074468016624451, 0.26749876141548157, 0.26749876141548157, 0.9915618896484375, 0.9915618896484375, 0.9999154806137085, 0.9999154806137085],
+                    [0.13673721253871918, 0.13673721253871918, 0.1699671596288681, 0.1699671596288681, 0.9902160167694092, 0.9902160167694092, 0.9999020099639893, 0.9999020099639893]
+                ]]]);
+        const expected_sine_cache = tf.tensor([[[
+                    [0, 0, 0, 0, 0, 0, 0, 0],
+                    [0.8414709568023682, 0.8414709568023682, 0.0998334214091301, 0.0998334214091301, 0.009999833069741726, 0.009999833069741726, 0.0009999999310821295, 0.0009999999310821295],
+                    [0.9092974066734314, 0.9092974066734314, 0.19866932928562164, 0.19866932928562164, 0.019998665899038315, 0.019998665899038315, 0.0019999986980110407, 0.0019999986980110407],
+                    [0.14112000167369843, 0.14112000167369843, 0.29552021622657776, 0.29552021622657776, 0.029995499178767204, 0.029995499178767204, 0.0029999956022948027, 0.0029999956022948027],
+                    [-0.756802499294281, -0.756802499294281, 0.3894183337688446, 0.3894183337688446, 0.03998933359980583, 0.03998933359980583, 0.003999989479780197, 0.003999989479780197],
+                    [-0.9589242935180664, -0.9589242935180664, 0.4794255495071411, 0.4794255495071411, 0.04997916519641876, 0.04997916519641876, 0.0049999793991446495, 0.0049999793991446495],
+                    [-0.279415488243103, -0.279415488243103, 0.5646424889564514, 0.5646424889564514, 0.059964004904031754, 0.059964004904031754, 0.0059999641962349415, 0.0059999641962349415],
+                    [0.6569865942001343, 0.6569865942001343, 0.6442176699638367, 0.6442176699638367, 0.06994284689426422, 0.06994284689426422, 0.0069999429397284985, 0.0069999429397284985],
+                    [0.9893582463264465, 0.9893582463264465, 0.7173560857772827, 0.7173560857772827, 0.07991468906402588, 0.07991468906402588, 0.007999914698302746, 0.007999914698302746],
+                    [0.41211849451065063, 0.41211849451065063, 0.7833269238471985, 0.7833269238471985, 0.08987854421138763, 0.08987854421138763, 0.008999879471957684, 0.008999879471957684],
+                    [-0.5440211296081543, -0.5440211296081543, 0.8414709568023682, 0.8414709568023682, 0.0998334139585495, 0.0998334139585495, 0.0099998340010643, 0.0099998340010643],
+                    [-0.9999902248382568, -0.9999902248382568, 0.8912073969841003, 0.8912073969841003, 0.10977829992771149, 0.10977829992771149, 0.010999779216945171, 0.010999779216945171],
+                    [-0.5365729331970215, -0.5365729331970215, 0.9320390820503235, 0.9320390820503235, 0.11971220374107361, 0.11971220374107361, 0.011999712325632572, 0.011999712325632572],
+                    [0.4201670289039612, 0.4201670289039612, 0.9635581970214844, 0.9635581970214844, 0.12963414192199707, 0.12963414192199707, 0.012999634258449078, 0.012999634258449078],
+                    [0.9906073808670044, 0.9906073808670044, 0.9854497313499451, 0.9854497313499451, 0.13954311609268188, 0.13954311609268188, 0.013999543152749538, 0.013999543152749538]
+                ]]]);
+        const [cosine_cache, sine_cache] = rope.getWeights();
+        expect(await cosine_cache?.sub(expected_cosine_cache).sum().array()).toBeLessThanOrEqual(1e-6);
+        expect(await sine_cache?.sub(expected_sine_cache).sum().array()).toBeLessThanOrEqual(1e-6);
+    });
+    test("rotate inputs", async () => {
+        const rope = new RotaryPositionEmbedding({ dim: 8, maxSequenceLength: 15 });
+        const x = tf.tensor([[[
+                    [0.0766048, 0.5706575, 0.6705932, 0.5273118, 0.4794086, 0.9378104, 0.9888024, 0.6926053],
+                    [0.9064133, 0.5875182, 0.1681865, 0.3833345, 0.9901192, 0.4677338, 0.3353315, 0.02699],
+                    [0.3033573, 0.4139377, 0.4062586, 0.9705839, 0.3582608, 0.328775, 0.1340587, 0.2193414],
+                    [0.5565202, 0.4334963, 0.9912352, 0.3388563, 0.7991487, 0.1911893, 0.1140554, 0.6949552]
+                ]]
+        ]); // batch=1, seq = 1, heads=4, embedDim=8
+        const expected_output = tf.tensor([[[
+                    [0.07660479843616486, 0.57065749168396, 0.6705932021141052, 0.5273118019104004, 0.4794085919857025, 0.9378104209899902, 0.9888023734092712, 0.6926053166389465],
+                    [-0.004642367362976074, 1.08015775680542, 0.12907665967941284, 0.39820998907089233, 0.9853923320770264, 0.47761136293411255, 0.33530429005622864, 0.027325313538312912],
+                    [-0.5026336908340454, 0.10358311235904694, 0.20533521473407745, 1.0319478511810303, 0.3516140580177307, 0.33587393164634705, 0.1336197406053543, 0.21960905194282532],
+                    [-0.6121258735656738, -0.3506217896938324, 0.8468242287635803, 0.6166517734527588, 0.7930541634559631, 0.2150741070508957, 0.11197001487016678, 0.695294201374054]
+                ]]]);
+        const output = rope.apply(x);
+        expect(await expected_output.sub(output).sum().array()).toBeLessThan(1e-6);
+        expect(rope.computeOutputShape(x.shape)).toEqual(x.shape);
+        expect(rope.computeOutputShape([x.shape])).toEqual(x.shape);
+    });
+    test("expand cache when input sequences are larger than rope's max sequence length", async () => {
+        const dim = 8;
+        const rope = new RotaryPositionEmbedding({ dim, maxSequenceLength: 15, theta: 1_000_000 });
+        const larger_sequence = 20;
+        const even_larger_sequence = 50;
+        rope.apply(tf.randomUniform([1, 1, larger_sequence, dim]));
+        rope.getWeights().forEach(weight => {
+            expect(weight.shape).toEqual([1, 1, 32, dim]);
+        });
+        rope.apply([tf.randomUniform([1, 1, even_larger_sequence, dim])]);
+        rope.getWeights().forEach(weight => {
+            expect(weight.shape).toEqual([1, 1, 64, dim]);
+        });
+    });
+    test("create layer", async () => {
+        // dim must be even
+        expect(() => new RotaryPositionEmbedding({ dim: 7, maxSequenceLength: 15 })).toThrow();
+        expect(() => new RotaryPositionEmbedding({ dim: 8, maxSequenceLength: 25 })).not.toThrow();
+    });
+});
+//# sourceMappingURL=rotary_position_embedding.test.js.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"rotary_position_embedding.test.js","sourceRoot":"","sources":["../../src/layers/rotary_position_embedding.test.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,uBAAuB,EAAE,MAAM,oCAAoC,CAAC;AAC7E,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAEvC,sDAAsD;AACtD,wEAAwE;AACxE,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,CAAC,SAAS,EAAE,KAAK,CAAC,CAAC;AAG/B,QAAQ,CAAC,+BAA+B,EAAE,GAAG,EAAE;IAC3C,IAAI,CAAC,cAAc,EAAE,KAAK,IAAI,EAAE;QAC5B,MAAM,IAAI,GAAG,IAAI,uBAAuB,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE,iBAAiB,EAAE,EAAE,EAAE,CAAC,CAAC;QAC5E,IAAI,CAAC,KAAK,CAAC,EAAE,CAAC,CAAC;QAEf,MAAM,qBAAqB,GAAG,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC;oBACtC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;oBACxB,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAChK,CAAC,CAAC,iBAAiB,EAAE,CAAC,iBAAiB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAChK,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAClK,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAClK,CAAC,mBAAmB,EAAE,mBAAmB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAClK,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,iBAAiB,EAAE,iBAAiB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAC9J,CAAC,eAAe,EAAE,eAAe,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAC1J,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAClK,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAClK,CAAC,CAAC,gBAAgB,EAAE,CAAC,gBAAgB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAC9J,CAAC,oBAAoB,EAAE,oBAAoB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,iBAAiB,EAAE,iBAAiB,CAAC;oBAClK,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAChK,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAClK,CAAC,mBAAmB,EAAE,mBAAmB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;iBACrK,CAAC,CAAC,CAAC,CAAC;QAEL,MAAM,mBAAmB,GAAG,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC;oBACpC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC;oBACxB,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,oBAAoB,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,qBAAqB,CAAC;oBAC1K,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,qBAAqB,CAAC;oBAC5K,CAAC,mBAAmB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,qBAAqB,CAAC;oBAC9K,CAAC,CAAC,iBAAiB,EAAE,CAAC,iBAAiB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,CAAC;oBACtK,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,qBAAqB,EAAE,qBAAqB,CAAC;oBAC1K,CAAC,CAAC,iBAAiB,EAAE,CAAC,iBAAiB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,oBAAoB,EAAE,oBAAoB,EAAE,qBAAqB,EAAE,qBAAqB,CAAC;oBAC1K,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,qBAAqB,EAAE,qBAAqB,CAAC;oBACxK,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,CAAC;oBACtK,CAAC,mBAAmB,EAAE,mBAAmB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,CAAC;oBACxK,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAClK,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,CAAC;oBACxK,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,CAAC;oBACxK,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,CAAC;oBACtK,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,EAAE,oBAAoB,CAAC;iBACzK,CAAC,CAAC,CAAC,CAAC;QAEL,MAAM,CAAC,YAAY,EAAE,UAAU,CAAC,GAAG,IAAI,CAAC,UAAU,EAAE,CAAC;QAErD,MAAM,CAAC,MAAM,YAAY,EAAE,GAAG,CAAC,qBAAqB,CAAC,CAAC,GAAG,EAAE,CAAC,KAAK,EAAY,CAAC,CAAC,mBAAmB,CAAC,IAAI,CAAC,CAAC;QACzG,MAAM,CAAC,MAAM,UAAU,EAAE,GAAG,CAAC,mBAAmB,CAAC,CAAC,GAAG,EAAE,CAAC,KAAK,EAAY,CAAC,CAAC,mBAAmB,CAAC,IAAI,CAAC,CAAC;IACzG,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,eAAe,EAAE,KAAK,IAAI,EAAE;QAC7B,MAAM,IAAI,GAAG,IAAI,uBAAuB,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE,iBAAiB,EAAE,EAAE,EAAE,CAAC,CAAC;QAE5E,MAAM,CAAC,GAAG,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC;oBAClB,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;oBACxF,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,OAAO,CAAC;oBACtF,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,QAAQ,EAAE,SAAS,EAAE,SAAS,CAAC;oBACvF,CAAC,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,EAAE,SAAS,CAAC;iBAAC,CAAC;SAC7F,CAAC,CAAC,CAAC,wCAAwC;QAE5C,MAAM,eAAe,GAAG,EAAE,CAAC,MAAM,CAAC,CAAC,CAAC;oBAChC,CAAC,mBAAmB,EAAE,gBAAgB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,CAAC;oBAC/J,CAAC,CAAC,oBAAoB,EAAE,gBAAgB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,oBAAoB,CAAC;oBACvK,CAAC,CAAC,kBAAkB,EAAE,mBAAmB,EAAE,mBAAmB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,kBAAkB,EAAE,mBAAmB,CAAC;oBACrK,CAAC,CAAC,kBAAkB,EAAE,CAAC,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,kBAAkB,EAAE,mBAAmB,EAAE,iBAAiB,CAAC;iBACrK,CAAC,CAAC,CAAC,CAAC;QAEL,MAAM,MAAM,GAAG,IAAI,CAAC,KAAK,CAAC,CAAC,CAAc,CAAC;QAE1C,MAAM,CAAC,MAAM,eAAe,CAAC,GAAG,CAAC,MAAM,CAAC,CAAC,GAAG,EAAE,CAAC,KAAK,EAAY,CAAC,CAAC,YAAY,CAAC,IAAI,CAAC,CAAC;QACrF,MAAM,CAAC,IAAI,CAAC,kBAAkB,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC;QAC1D,MAAM,CAAC,IAAI,CAAC,kBAAkB,CAAC,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,KAAK,CAAC,CAAC;IAChE,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,8EAA8E,EAAE,KAAK,IAAI,EAAE;QAC5F,MAAM,GAAG,GAAG,CAAC,CAAC;QACd,MAAM,IAAI,GAAG,IAAI,uBAAuB,CAAC,EAAE,GAAG,EAAE,iBAAiB,EAAE,EAAE,EAAE,KAAK,EAAE,SAAS,EAAE,CAAC,CAAC;QAC3F,MAAM,eAAe,GAAG,EAAE,CAAC;QAC3B,MAAM,oBAAoB,GAAG,EAAE,CAAC;QAEhC,IAAI,CAAC,KAAK,CAAC,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,eAAe,EAAE,GAAG,CAAC,CAAC,CAAC,CAAC;QAE3D,IAAI,CAAC,UAAU,EAAE,CAAC,OAAO,CAAC,MAAM,CAAC,EAAE;YAC/B,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,GAAG,CAAC,CAAC,CAAC;QAClD,CAAC,CAAC,CAAC;QAEH,IAAI,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,aAAa,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,oBAAoB,EAAE,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC;QAElE,IAAI,CAAC,UAAU,EAAE,CAAC,OAAO,CAAC,MAAM,CAAC,EAAE;YAC/B,MAAM,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,EAAE,GAAG,CAAC,CAAC,CAAC;QAClD,CAAC,CAAC,CAAC;IACP,CAAC,CAAC,CAAA;IAGF,IAAI,CAAC,cAAc,EAAE,KAAK,IAAI,EAAE;QAC5B,mBAAmB;QACnB,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,uBAAuB,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE,iBAAiB,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QACvF,MAAM,CAAC,GAAG,EAAE,CAAC,IAAI,uBAAuB,CAAC,EAAE,GAAG,EAAE,CAAC,EAAE,iBAAiB,EAAE,EAAE,EAAE,CAAC,CAAC,CAAC,GAAG,CAAC,OAAO,EAAE,CAAC;IAC/F,CAAC,CAAC,CAAA;AACN,CAAC,CAAC,CAAC"}

package/dist/layers/token_and_positional_embedding.d.ts

@@ -0,0 +1,47 @@
+import * as tf from '@tensorflow/tfjs';
+import { type LayerArgs } from '@tensorflow/tfjs-layers/dist/engine/topology';
+import { type Kwargs } from '@tensorflow/tfjs-layers/dist/types';
+import { type PositionalEncodingArgs } from '../layers/positional_encoding';
+export interface TokenAndPositionalEmbeddingArgs extends LayerArgs, PositionalEncodingArgs {
+    vocabularySize: number;
+    dropout?: number;
+}
+/**
+ * This class implements combines sinusoidal positional encoding from the
+ * 2017 paper "Attention Is All You Need" with a normal embedding layer to
+ * form a simplified single embedding layer.
+ *
+ * This layer accepts tokenized inputs of the shape `[ batch, tokens ]` and runs
+ * it through an embedding layer before adding sinusoidal positional encoding.
+ *
+ * @param embedDim the size of each token/word's embedding
+ * @param vocabularySize the number of tokens to embed
+ * @param maxSequenceLength the max number of tokens (words) per input (sentence), default `5120`
+ * @param dropout applies dropout to the positionally encoded embeddings, default `0.1`
+ */
+export declare class TokenAndPositionalEmbedding extends tf.layers.Layer {
+    static className: string;
+    private readonly embedDim;
+    private readonly vocabularySize;
+    private embedding;
+    private positional;
+    private readonly maxSequenceLength;
+    private readonly dropout;
+    private dropoutLayer;
+    constructor({ embedDim, vocabularySize, maxSequenceLength, dropout, ...args }: TokenAndPositionalEmbeddingArgs);
+    /**
+     * Forward propagation.
+     */
+    call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor<tf.Rank>;
+    /**
+     * Build the sublayers and enable serialization
+     */
+    build(inputShape: tf.Shape | tf.Shape[]): void;
+    /**
+     * The output shape, for an input shape of [batch, sequences], is
+     * [batch, sequences, embedDim]
+     */
+    computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[];
+    getConfig(): tf.serialization.ConfigDict;
+}
+//# sourceMappingURL=token_and_positional_embedding.d.ts.map

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

@@ -0,0 +1 @@
+{"version":3,"file":"token_and_positional_embedding.d.ts","sourceRoot":"","sources":["../../src/layers/token_and_positional_embedding.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;AAEjE,OAAO,EAAsB,KAAK,sBAAsB,EAAE,MAAM,+BAA+B,CAAC;AAGhG,MAAM,WAAW,+BAAgC,SAAQ,SAAS,EAAE,sBAAsB;IACtF,cAAc,EAAE,MAAM,CAAC;IACvB,OAAO,CAAC,EAAE,MAAM,CAAA;CACnB;AAGD;;;;;;;;;;;;GAYG;AACH,qBAAa,2BAA4B,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IAC5D,MAAM,CAAC,SAAS,SAAiC;IAEjD,OAAO,CAAC,QAAQ,CAAC,QAAQ,CAAS;IAClC,OAAO,CAAC,QAAQ,CAAC,cAAc,CAAS;IACxC,OAAO,CAAC,SAAS,CAAkB;IAEnC,OAAO,CAAC,UAAU,CAAiB;IACnC,OAAO,CAAC,QAAQ,CAAC,iBAAiB,CAAS;IAC3C,OAAO,CAAC,QAAQ,CAAC,OAAO,CAAS;IAEjC,OAAO,CAAC,YAAY,CAAkB;gBAG1B,EAAE,QAAQ,EAAE,cAAc,EAAE,iBAAiB,EAAE,OAAO,EAAE,GAAG,IAAI,EAAE,EAAE,+BAA+B;IA0B9G;;OAEG;IACM,IAAI,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE,EAAE,MAAM,EAAE,MAAM;IAe7D;;OAEG;IACM,KAAK,CAAC,UAAU,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,IAAI;IAgCvD;;;OAGG;IACM,kBAAkB,CAAC,UAAU,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE;IAQ5E,SAAS,IAAI,EAAE,CAAC,aAAa,CAAC,UAAU;CAcpD"}