@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

package/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 rkuang9
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

package/README.md

@@ -0,0 +1,47 @@
+# tfjs-stellar
+An extension of TensorFlow.js for implementing large language models.
+
+
+# Layers
+- MultiHeadAttention
+- CachedRopeMultiHeadAttention
+- TransformerDecoder
+- TransformerEncoder
+- GPT2DecoderBlock
+- RotaryPositionEmbedding
+- PositionalEncoding
+- TokenAndPositionalEmbedding
+
+> **Warning**:
+> These layers are not one-to-one replications of the TensorFlow Keras equivalents
+
+
+## Models
+- LlmModel
+- GptModel
+- KvCache
+- UNetModel
+
+
+## Masks
+- Causal
+- Packing
+
+
+## Example
+
+```ts
+import * as tfs from "@stellarapp/tfjs-stellar";
+import * as tf from "@tensorflow/tfjs";
+
+const attention = tfs.layers.multiheadAttention({ numHeads: 1, embedDim: 64 });
+const output = attention.apply(tf.randomUniform([1, 5, 64]));
+
+const gpt_model = tfs.models.gptModel({ numLayers: 1, numHeads: 1, embedDim: 64, vocabSize: 128 });
+gpt_model.compile({ loss: "sparseCategoricalCrossentropy", optimizer: "adam" });
+gpt_model.summary();
+
+// see https://js.tensorflow.org/api/latest/#data.generator
+// on how to create a generator dataset
+//gpt_model.fitDataset(your_generator_dataset, { epochs: 1 });
+```

package/dist/index.d.ts

@@ -0,0 +1,7 @@
+export * as layers from "./layers";
+export * as models from "./models";
+export * as losses from "./losses";
+export * as masks from "./masks";
+export { KvCache as kvCache, KvCacheContainer as kvCacheContainer } from "./kv_cache";
+export * as metrics from "./metrics";
+//# sourceMappingURL=index.d.ts.map

package/dist/index.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,MAAM,MAAM,UAAU,CAAC;AACnC,OAAO,KAAK,MAAM,MAAM,UAAU,CAAC;AACnC,OAAO,KAAK,MAAM,MAAM,UAAU,CAAC;AAEnC,OAAO,KAAK,KAAK,MAAM,SAAS,CAAC;AACjC,OAAO,EAAE,OAAO,IAAI,OAAO,EAAE,gBAAgB,IAAI,gBAAgB,EAAE,MAAM,YAAY,CAAC;AACtF,OAAO,KAAK,OAAO,MAAM,WAAW,CAAC"}

package/dist/index.js

@@ -0,0 +1,7 @@
+export * as layers from "./layers";
+export * as models from "./models";
+export * as losses from "./losses";
+export * as masks from "./masks";
+export { KvCache as kvCache, KvCacheContainer as kvCacheContainer } from "./kv_cache";
+export * as metrics from "./metrics";
+//# sourceMappingURL=index.js.map

package/dist/index.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"index.js","sourceRoot":"","sources":["../src/index.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,MAAM,MAAM,UAAU,CAAC;AACnC,OAAO,KAAK,MAAM,MAAM,UAAU,CAAC;AACnC,OAAO,KAAK,MAAM,MAAM,UAAU,CAAC;AAEnC,OAAO,KAAK,KAAK,MAAM,SAAS,CAAC;AACjC,OAAO,EAAE,OAAO,IAAI,OAAO,EAAE,gBAAgB,IAAI,gBAAgB,EAAE,MAAM,YAAY,CAAC;AACtF,OAAO,KAAK,OAAO,MAAM,WAAW,CAAC"}

package/dist/jest.config.d.ts

@@ -0,0 +1,8 @@
+/**
+ * For a detailed explanation regarding each configuration property, visit:
+ * https://jestjs.io/docs/configuration
+ */
+import type { Config } from 'jest';
+declare const config: Config;
+export default config;
+//# sourceMappingURL=jest.config.d.ts.map

package/dist/jest.config.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"jest.config.d.ts","sourceRoot":"","sources":["../jest.config.ts"],"names":[],"mappings":"AAAA;;;GAGG;AAEH,OAAO,KAAK,EAAE,MAAM,EAAE,MAAM,MAAM,CAAC;AAEnC,QAAA,MAAM,MAAM,EAAE,MAiMb,CAAC;AAEF,eAAe,MAAM,CAAC"}

package/{jest.config.ts → dist/jest.config.js}

@@ -2,60 +2,43 @@
  * For a detailed explanation regarding each configuration property, visit:
  * https://jestjs.io/docs/configuration
  */
-
-import type { Config } from 'jest';
-
-const config: Config = {
+const config = {
     setupFiles: [],
-
     extensionsToTreatAsEsm: [".ts"],
-
     // A map from regular expressions to paths to transformers
     transform: {
         "^.+\.ts?$": ["ts-jest", {
-            useESM: true
-        }],
+                useESM: true
+            }],
     },
-
     // An array of regexp pattern strings that are matched against all test paths, matched tests are skipped
-    testPathIgnorePatterns: [
-        "/node_modules/", "e2e"
-    ],
-
+    // testPathIgnorePatterns: [
+    //     
+    // ],
     // A map from regular expressions to module names or to arrays of module names that allow to stub out resources with a single module
     moduleNameMapper: {
         "^@/(.*$)": "<rootDir>/src/$1"
     },
-
     // All imported modules in your tests should be mocked automatically
     // automock: false,
-
     // Stop running tests after `n` failures
     // bail: 0,
-
     // The directory where Jest should store its cached dependency information
     // cacheDirectory: "/private/var/folders/8x/0jgq0fqx5qzgtm1zc8xtrdc80000gn/T/jest_dx",
-
     // Automatically clear mock calls, instances, contexts and results before every test
     clearMocks: true,
-
     // Indicates whether the coverage information should be collected while executing the test
     collectCoverage: false,
-
     // An array of glob patterns indicating a set of files for which coverage information should be collected
     // collectCoverageFrom: undefined,
-
     // The directory where Jest should output its coverage files
     coverageDirectory: "coverage",
-
     // An array of regexp pattern strings used to skip coverage collection
     // coveragePathIgnorePatterns: [
     //   "/node_modules/"
     // ],
-
     // Indicates which provider should be used to instrument code for coverage
     // coverageProvider: "babel",
-
     // A list of reporter names that Jest uses when writing coverage reports
     // coverageReporters: [
     //   "json",
@@ -63,36 +46,26 @@ const config: Config = {
     //   "lcov",
     //   "clover"
     // ],
-
     // An object that configures minimum threshold enforcement for coverage results
     // coverageThreshold: undefined,
-
     // A path to a custom dependency extractor
     // dependencyExtractor: undefined,
-
     // Make calling deprecated APIs throw helpful error messages
     // errorOnDeprecated: false,
-
     // The default configuration for fake timers
     // fakeTimers: {
     //   "enableGlobally": false
     // },
-
     // Force coverage collection from ignored files using an array of glob patterns
     // forceCoverageMatch: [],
-
     // A path to a module which exports an async function that is triggered once before all test suites
     // globalSetup: undefined,
-
     // A path to a module which exports an async function that is triggered once after all test suites
     // globalTeardown: undefined,
-
     // A set of global variables that need to be available in all test environments
     // globals: {},
-
     // The maximum amount of workers used to run your tests. Can be specified as % or a number. E.g. maxWorkers: 10% will use 10% of your CPU amount + 1 as the maximum worker number. maxWorkers: 2 will use a maximum of 2 workers.
     // maxWorkers: "50%",
-
     // An array of file extensions your modules use
     // moduleFileExtensions: [
     //   "js",
@@ -104,100 +77,71 @@ const config: Config = {
     //   "json",
     //   "node"
     // ],
-
     // An array of regexp pattern strings, matched against all module paths before considered 'visible' to the module loader
-    // modulePathIgnorePatterns: [],
-
+    modulePathIgnorePatterns: ["<rootDir>/dist/"],
     // Activates notifications for test results
     // notify: false,
-
     // An enum that specifies notification mode. Requires { notify: true }
     // notifyMode: "failure-change",
-
     // A preset that is used as a base for Jest's configuration
     // preset: undefined,
-
     // Run tests from one or more projects
     // projects: undefined,
-
     // Use this configuration option to add custom reporters to Jest
     // reporters: undefined,
-
     // Automatically reset mock state before every test
     // resetMocks: false,
-
     // Reset the module registry before running each individual test
     // resetModules: false,
-
     // A path to a custom resolver
     // resolver: undefined,
-
     // Automatically restore mock state and implementation before every test
     // restoreMocks: false,
-
     // The root directory that Jest should scan for tests and modules within
     // rootDir: undefined,
-
     // A list of paths to directories that Jest should use to search for files in
     // roots: [
     //   "<rootDir>"
     // ],
-
     // Allows you to use a custom runner instead of Jest's default test runner
     // runner: "jest-runner",
-
     // The paths to modules that run some code to configure or set up the testing environment before each test
-
     // A list of paths to modules that run some code to configure or set up the testing framework before each test
     // setupFilesAfterEnv: [],
-
     // The number of seconds after which a test is considered as slow and reported as such in the results.
     // slowTestThreshold: 5,
-
     // A list of paths to snapshot serializer modules Jest should use for snapshot testing
     // snapshotSerializers: [],
-
     // The test environment that will be used for testing
     testEnvironment: "node",
-
     // Options that will be passed to the testEnvironment
     // testEnvironmentOptions: {},
-
     // Adds a location field to test results
     // testLocationInResults: false,
-
     // The glob patterns Jest uses to detect test files
     // testMatch: [
     //   "**/__tests__/**/*.[jt]s?(x)",
     //   "**/?(*.)+(spec|test).[tj]s?(x)"
     // ],
-
     // The regexp pattern or array of patterns that Jest uses to detect test files
     // testRegex: [],
-
     // This option allows the use of a custom results processor
     // testResultsProcessor: undefined,
-
     // This option allows use of a custom test runner
     // testRunner: "jest-circus/runner",
-
     // An array of regexp pattern strings that are matched against all source file paths, matched files will skip transformation
     // transformIgnorePatterns: [
     //   "/node_modules/",
     //   "\\.pnp\\.[^\\/]+$"
     // ],
-
     // An array of regexp pattern strings that are matched against all modules before the module loader will automatically return a mock for them
     // unmockedModulePathPatterns: undefined,
-
     // Indicates whether each individual test should be reported during the run
     // verbose: undefined,
-
     // An array of regexp patterns that are matched against all source file paths before re-running tests in watch mode
     // watchPathIgnorePatterns: [],
-
     // Whether to use watchman for file crawling
     // watchman: true,
 };
-
 export default config;
+//# sourceMappingURL=jest.config.js.map

package/dist/jest.config.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"jest.config.js","sourceRoot":"","sources":["../jest.config.ts"],"names":[],"mappings":"AAAA;;;GAGG;AAIH,MAAM,MAAM,GAAW;IACnB,UAAU,EAAE,EAAE;IAEd,sBAAsB,EAAE,CAAC,KAAK,CAAC;IAE/B,0DAA0D;IAC1D,SAAS,EAAE;QACP,WAAW,EAAE,CAAC,SAAS,EAAE;gBACrB,MAAM,EAAE,IAAI;aACf,CAAC;KACL;IAED,wGAAwG;IACxG,4BAA4B;IAC5B,OAAO;IACP,KAAK;IAEL,oIAAoI;IACpI,gBAAgB,EAAE;QACd,UAAU,EAAE,kBAAkB;KACjC;IAED,oEAAoE;IACpE,mBAAmB;IAEnB,wCAAwC;IACxC,WAAW;IAEX,0EAA0E;IAC1E,sFAAsF;IAEtF,oFAAoF;IACpF,UAAU,EAAE,IAAI;IAEhB,0FAA0F;IAC1F,eAAe,EAAE,KAAK;IAEtB,yGAAyG;IACzG,kCAAkC;IAElC,4DAA4D;IAC5D,iBAAiB,EAAE,UAAU;IAE7B,sEAAsE;IACtE,gCAAgC;IAChC,qBAAqB;IACrB,KAAK;IAEL,0EAA0E;IAC1E,6BAA6B;IAE7B,wEAAwE;IACxE,uBAAuB;IACvB,YAAY;IACZ,YAAY;IACZ,YAAY;IACZ,aAAa;IACb,KAAK;IAEL,+EAA+E;IAC/E,gCAAgC;IAEhC,0CAA0C;IAC1C,kCAAkC;IAElC,4DAA4D;IAC5D,4BAA4B;IAE5B,4CAA4C;IAC5C,gBAAgB;IAChB,4BAA4B;IAC5B,KAAK;IAEL,+EAA+E;IAC/E,0BAA0B;IAE1B,mGAAmG;IACnG,0BAA0B;IAE1B,kGAAkG;IAClG,6BAA6B;IAE7B,+EAA+E;IAC/E,eAAe;IAEf,iOAAiO;IACjO,qBAAqB;IAErB,+CAA+C;IAC/C,0BAA0B;IAC1B,UAAU;IACV,WAAW;IACX,WAAW;IACX,WAAW;IACX,UAAU;IACV,WAAW;IACX,YAAY;IACZ,WAAW;IACX,KAAK;IAEL,wHAAwH;IACxH,wBAAwB,EAAE,CAAC,iBAAiB,CAAC;IAE7C,2CAA2C;IAC3C,iBAAiB;IAEjB,sEAAsE;IACtE,gCAAgC;IAEhC,2DAA2D;IAC3D,qBAAqB;IAErB,sCAAsC;IACtC,uBAAuB;IAEvB,gEAAgE;IAChE,wBAAwB;IAExB,mDAAmD;IACnD,qBAAqB;IAErB,gEAAgE;IAChE,uBAAuB;IAEvB,8BAA8B;IAC9B,uBAAuB;IAEvB,wEAAwE;IACxE,uBAAuB;IAEvB,wEAAwE;IACxE,sBAAsB;IAEtB,6EAA6E;IAC7E,WAAW;IACX,gBAAgB;IAChB,KAAK;IAEL,0EAA0E;IAC1E,yBAAyB;IAEzB,0GAA0G;IAE1G,8GAA8G;IAC9G,0BAA0B;IAE1B,sGAAsG;IACtG,wBAAwB;IAExB,sFAAsF;IACtF,2BAA2B;IAE3B,qDAAqD;IACrD,eAAe,EAAE,MAAM;IAEvB,qDAAqD;IACrD,8BAA8B;IAE9B,wCAAwC;IACxC,gCAAgC;IAEhC,mDAAmD;IACnD,eAAe;IACf,mCAAmC;IACnC,qCAAqC;IACrC,KAAK;IAEL,8EAA8E;IAC9E,iBAAiB;IAEjB,2DAA2D;IAC3D,mCAAmC;IAEnC,iDAAiD;IACjD,oCAAoC;IAEpC,4HAA4H;IAC5H,6BAA6B;IAC7B,sBAAsB;IACtB,wBAAwB;IACxB,KAAK;IAEL,6IAA6I;IAC7I,yCAAyC;IAEzC,2EAA2E;IAC3E,sBAAsB;IAEtB,mHAAmH;IACnH,+BAA+B;IAE/B,4CAA4C;IAC5C,kBAAkB;CACrB,CAAC;AAEF,eAAe,MAAM,CAAC"}

package/dist/kv_cache.d.ts

@@ -0,0 +1,53 @@
+import * as tf from "@tensorflow/tfjs";
+export interface KvCacheArgs {
+    batchSize: number;
+    maxSequenceLength: number;
+    numHeads: number;
+    headDim: number;
+    dtype?: tf.DataType;
+}
+/**
+ * A container for KV caches. A model should initialize one KV cache
+ */
+export declare class KvCacheContainer {
+    protected caches: Map<string, KvCache>;
+    protected max_sequence_length: number;
+    constructor(maxSequenceLength: number);
+    create(id: string, args: Omit<KvCacheArgs, "maxSequenceLength">): void;
+    /**
+     * The key and value tensors should have the shape (post head split, etc) `[batch, heads, seq, head_dim]`
+     */
+    update(id: string, key: tf.Tensor4D, value: tf.Tensor4D): {
+        keyCache: tf.Variable<tf.Rank.R4>;
+        valueCache: tf.Variable<tf.Rank.R4>;
+    } | undefined;
+    reset(): void;
+    dispose(): void;
+    get size(): number;
+    get maxSequenceLength(): number;
+}
+export declare class KvCache {
+    protected key_cache: tf.Variable<tf.Rank.R4>;
+    protected value_cache: tf.Variable<tf.Rank.R4>;
+    protected current_position: number;
+    protected batch_size: number;
+    protected max_sequence_length: number;
+    protected num_kv_heads: number;
+    protected head_dim: number;
+    constructor({ batchSize, maxSequenceLength, numHeads, headDim, dtype }: KvCacheArgs);
+    /**
+     * The key and value tensors should have the shape (post head split, etc) `[batch, heads, seq, head_dim]`
+     */
+    update(key: tf.Tensor4D, value: tf.Tensor4D): {
+        keyCache: tf.Variable<tf.Rank.R4>;
+        valueCache: tf.Variable<tf.Rank.R4>;
+    };
+    protected mergeIntoCache(new_value: tf.Tensor4D, current_cache: tf.Tensor4D): tf.Tensor4D;
+    reset(): void;
+    dispose(): void;
+    /**
+     * The size of the KV cache, also the number of tokens since the first one.
+     */
+    get size(): number;
+}
+//# sourceMappingURL=kv_cache.d.ts.map

package/dist/kv_cache.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"kv_cache.d.ts","sourceRoot":"","sources":["../src/kv_cache.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAGvC,MAAM,WAAW,WAAW;IACxB,SAAS,EAAE,MAAM,CAAC;IAClB,iBAAiB,EAAE,MAAM,CAAC;IAC1B,QAAQ,EAAE,MAAM,CAAC;IACjB,OAAO,EAAE,MAAM,CAAC;IAChB,KAAK,CAAC,EAAE,EAAE,CAAC,QAAQ,CAAA;CACtB;AAGD;;GAEG;AACH,qBAAa,gBAAgB;IACzB,SAAS,CAAC,MAAM,uBAA8B;IAC9C,SAAS,CAAC,mBAAmB,EAAE,MAAM,CAAC;gBAG1B,iBAAiB,EAAE,MAAM;IAS9B,MAAM,CAAC,EAAE,EAAE,MAAM,EAAE,IAAI,EAAE,IAAI,CAAC,WAAW,EAAE,mBAAmB,CAAC;IAUtE;;OAEG;IACI,MAAM,CAAC,EAAE,EAAE,MAAM,EAAE,GAAG,EAAE,EAAE,CAAC,QAAQ,EAAE,KAAK,EAAE,EAAE,CAAC,QAAQ;;;;IA4BvD,KAAK;IAOL,OAAO;IAOd,IAAW,IAAI,WAGd;IAGD,IAAW,iBAAiB,WAE3B;CACJ;AAGD,qBAAa,OAAO;IAEhB,SAAS,CAAC,SAAS,EAAE,EAAE,CAAC,QAAQ,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;IAC7C,SAAS,CAAC,WAAW,EAAE,EAAE,CAAC,QAAQ,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,CAAA;IAG9C,SAAS,CAAC,gBAAgB,EAAE,MAAM,CAAK;IAEvC,SAAS,CAAC,UAAU,EAAE,MAAM,CAAC;IAC7B,SAAS,CAAC,mBAAmB,EAAE,MAAM,CAAC;IACtC,SAAS,CAAC,YAAY,EAAE,MAAM,CAAC;IAC/B,SAAS,CAAC,QAAQ,EAAE,MAAM,CAAC;gBAEf,EAAE,SAAS,EAAE,iBAAiB,EAAE,QAAQ,EAAE,OAAO,EAAE,KAAiB,EAAE,EAAE,WAAW;IAa/F;;OAEG;IACI,MAAM,CAAC,GAAG,EAAE,EAAE,CAAC,QAAQ,EAAE,KAAK,EAAE,EAAE,CAAC,QAAQ;;;;IAgClD,SAAS,CAAC,cAAc,CAAC,SAAS,EAAE,EAAE,CAAC,QAAQ,EAAE,aAAa,EAAE,EAAE,CAAC,QAAQ;IAqBpE,KAAK,IAAI,IAAI;IAab,OAAO,IAAI,IAAI;IAMtB;;OAEG;IACH,IAAI,IAAI,IAAI,MAAM,CAEjB;CAEJ"}

package/{src/kv_cache.ts → dist/kv_cache.js}

@@ -1,205 +1,135 @@
 import * as tf from "@tensorflow/tfjs";
-
-
-export interface KvCacheArgs {
-    batchSize: number;
-    maxSequenceLength: number;
-    numHeads: number;
-    headDim: number;
-    dtype?: tf.DataType
-}
-
-
 /**
  * A container for KV caches. A model should initialize one KV cache
  */
 export class KvCacheContainer {
-    protected caches = new Map<string, KvCache>();
-    protected max_sequence_length: number;
-
-
-    constructor(maxSequenceLength: number) {
+    caches = new Map();
+    max_sequence_length;
+    constructor(maxSequenceLength) {
         if (!maxSequenceLength) {
             throw Error(`KvCacheContainer: expected KV cache maximum sequence length to be greater than 0, got: ${String(maxSequenceLength)}`);
         }
-
         this.max_sequence_length = maxSequenceLength;
     }
-
-
-    public create(id: string, args: Omit<KvCacheArgs, "maxSequenceLength">) {
+    create(id, args) {
         const new_cache = new KvCache({
             ...args,
             maxSequenceLength: this.max_sequence_length
         });
-
         this.caches.set(id, new_cache);
     }
-
-
     /**
      * The key and value tensors should have the shape (post head split, etc) `[batch, heads, seq, head_dim]`
      */
-    public update(id: string, key: tf.Tensor4D, value: tf.Tensor4D) {
+    update(id, key, value) {
         const kv_cache = this.caches.get(id);
-
         if (!kv_cache) {
             return undefined;
         }
-
         const { keyCache, valueCache } = kv_cache.update(key, value);
-
         // slicing to get only the past key and value projections, but normally
         // in TensorFlow and PyTorch the full cache is returned and masked for
         // graph purposes
         return tf.tidy(() => {
-            const k_cache = keyCache.slice(
-                [0, 0, 0, 0],
-                [keyCache.shape[0], keyCache.shape[1], kv_cache.size, keyCache.shape[3]]);
-            const v_cache = valueCache.slice(
-                [0, 0, 0, 0],
-                [valueCache.shape[0], valueCache.shape[1], kv_cache.size, valueCache.shape[3]]);
-
+            const k_cache = keyCache.slice([0, 0, 0, 0], [keyCache.shape[0], keyCache.shape[1], kv_cache.size, keyCache.shape[3]]);
+            const v_cache = valueCache.slice([0, 0, 0, 0], [valueCache.shape[0], valueCache.shape[1], kv_cache.size, valueCache.shape[3]]);
             return {
                 keyCache: k_cache,
                 valueCache: v_cache
-            }
-        })
+            };
+        });
     }
-
-
-    public reset() {
+    reset() {
         this.caches.forEach(cache => {
             cache.reset();
-        })
+        });
     }
-
-
-    public dispose() {
+    dispose() {
         this.caches.forEach(cache => {
             cache.dispose();
-        })
+        });
     }
-
-
-    public get size() {
+    get size() {
         // the size of all KV caches are expected to be the same, just use the first one
         return this.caches.entries().next().value?.[1].size ?? 0;
     }
-
-
-    public get maxSequenceLength() {
+    get maxSequenceLength() {
         return this.max_sequence_length;
     }
 }
-
-
 export class KvCache {
-
-    protected key_cache: tf.Variable<tf.Rank.R4>;
-    protected value_cache: tf.Variable<tf.Rank.R4>
-
+    key_cache;
+    value_cache;
     // the size of the KV cache, represents the number of tokens since the first chat token
-    protected current_position: number = 0;
-
-    protected batch_size: number;
-    protected max_sequence_length: number;
-    protected num_kv_heads: number;
-    protected head_dim: number;
-
-    constructor({ batchSize, maxSequenceLength, numHeads, headDim, dtype = "float32" }: KvCacheArgs) {
-        const cache_shape = [batchSize, numHeads, maxSequenceLength, headDim] as [number, number, number, number];
-
+    current_position = 0;
+    batch_size;
+    max_sequence_length;
+    num_kv_heads;
+    head_dim;
+    constructor({ batchSize, maxSequenceLength, numHeads, headDim, dtype = "float32" }) {
+        const cache_shape = [batchSize, numHeads, maxSequenceLength, headDim];
         this.key_cache = tf.variable(tf.zeros(cache_shape, dtype), false);
         this.value_cache = tf.variable(tf.zeros(cache_shape, dtype), false);
-
         this.batch_size = batchSize;
         this.max_sequence_length = maxSequenceLength;
         this.num_kv_heads = numHeads;
         this.head_dim = headDim;
     }
-
-
     /**
      * The key and value tensors should have the shape (post head split, etc) `[batch, heads, seq, head_dim]`
      */
-    public update(key: tf.Tensor4D, value: tf.Tensor4D) {
+    update(key, value) {
         const batch_size = key.shape[0];
         const seq_len = key.shape[2];
-
         if (batch_size > this.key_cache.shape[0]) {
             throw Error(`The current KV cache has been set up with a batch size of` +
-                ` ${this.key_cache.shape[0]}, but found new key tensors with batch size ${batch_size}`)
+                ` ${this.key_cache.shape[0]}, but found new key tensors with batch size ${batch_size}`);
         }
-
         if (this.current_position + seq_len > this.max_sequence_length) {
             throw Error(`The KV cache has exceeded its maximum sequence length of ${this.max_sequence_length}. Use a larger value.`);
         }
-
         const new_key_cache = this.mergeIntoCache(key, this.key_cache);
         const new_value_cache = this.mergeIntoCache(value, this.value_cache);
-
         this.key_cache.assign(new_key_cache);
         this.value_cache.assign(new_value_cache);
-
         new_key_cache.dispose();
         new_value_cache.dispose();
-
         // advance the pointer to reflect the updated cache's current
         this.current_position += seq_len;
-
         return {
             keyCache: this.key_cache,
             valueCache: this.value_cache,
-        }
+        };
     }
-
-
-    protected mergeIntoCache(new_value: tf.Tensor4D, current_cache: tf.Tensor4D) {
+    mergeIntoCache(new_value, current_cache) {
         const seq_len = new_value.shape[2];
-
         return tf.tidy(() => {
-
-            const historical = current_cache.slice(
-                [0, 0, 0, 0],
-                [this.batch_size, this.num_kv_heads, this.current_position, this.head_dim]);
-
-            const future = current_cache.slice(
-                [0, 0, this.current_position + seq_len, 0],
-                [this.batch_size, this.num_kv_heads, this.max_sequence_length - this.current_position - seq_len, this.head_dim]);
-
+            const historical = current_cache.slice([0, 0, 0, 0], [this.batch_size, this.num_kv_heads, this.current_position, this.head_dim]);
+            const future = current_cache.slice([0, 0, this.current_position + seq_len, 0], [this.batch_size, this.num_kv_heads, this.max_sequence_length - this.current_position - seq_len, this.head_dim]);
             // merge the new tensor into the current cache to create a new, larger, cache,
             // this is different from Python immplementations because TFJS tensors are immutable,
             // because we cannot update a slice, we must slice and concat
             return tf.concat([historical, new_value, future], 2);
-        })
+        });
     }
-
-
-    public reset(): void {
+    reset() {
         this.current_position = 0;
-
         tf.tidy(() => {
             const key_cache_shape = this.key_cache.shape;
             const value_cache_shape = this.value_cache.shape;
-
             this.key_cache.assign(tf.zeros(key_cache_shape));
             this.value_cache.assign(tf.zeros(value_cache_shape));
         });
     }
-
-
-    public dispose(): void {
+    dispose() {
         this.key_cache.dispose();
         this.value_cache.dispose();
     }
-
-
     /**
      * The size of the KV cache, also the number of tokens since the first one.
      */
-    get size(): number {
+    get size() {
         return this.current_position;
     }
-
 }
+//# sourceMappingURL=kv_cache.js.map

package/dist/kv_cache.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"kv_cache.js","sourceRoot":"","sources":["../src/kv_cache.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAYvC;;GAEG;AACH,MAAM,OAAO,gBAAgB;IACf,MAAM,GAAG,IAAI,GAAG,EAAmB,CAAC;IACpC,mBAAmB,CAAS;IAGtC,YAAY,iBAAyB;QACjC,IAAI,CAAC,iBAAiB,EAAE,CAAC;YACrB,MAAM,KAAK,CAAC,0FAA0F,MAAM,CAAC,iBAAiB,CAAC,EAAE,CAAC,CAAC;QACvI,CAAC;QAED,IAAI,CAAC,mBAAmB,GAAG,iBAAiB,CAAC;IACjD,CAAC;IAGM,MAAM,CAAC,EAAU,EAAE,IAA4C;QAClE,MAAM,SAAS,GAAG,IAAI,OAAO,CAAC;YAC1B,GAAG,IAAI;YACP,iBAAiB,EAAE,IAAI,CAAC,mBAAmB;SAC9C,CAAC,CAAC;QAEH,IAAI,CAAC,MAAM,CAAC,GAAG,CAAC,EAAE,EAAE,SAAS,CAAC,CAAC;IACnC,CAAC;IAGD;;OAEG;IACI,MAAM,CAAC,EAAU,EAAE,GAAgB,EAAE,KAAkB;QAC1D,MAAM,QAAQ,GAAG,IAAI,CAAC,MAAM,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC;QAErC,IAAI,CAAC,QAAQ,EAAE,CAAC;YACZ,OAAO,SAAS,CAAC;QACrB,CAAC;QAED,MAAM,EAAE,QAAQ,EAAE,UAAU,EAAE,GAAG,QAAQ,CAAC,MAAM,CAAC,GAAG,EAAE,KAAK,CAAC,CAAC;QAE7D,uEAAuE;QACvE,sEAAsE;QACtE,iBAAiB;QACjB,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,OAAO,GAAG,QAAQ,CAAC,KAAK,CAC1B,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EACZ,CAAC,QAAQ,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,QAAQ,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,QAAQ,CAAC,IAAI,EAAE,QAAQ,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YAC9E,MAAM,OAAO,GAAG,UAAU,CAAC,KAAK,CAC5B,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EACZ,CAAC,UAAU,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,UAAU,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,QAAQ,CAAC,IAAI,EAAE,UAAU,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YAEpF,OAAO;gBACH,QAAQ,EAAE,OAAO;gBACjB,UAAU,EAAE,OAAO;aACtB,CAAA;QACL,CAAC,CAAC,CAAA;IACN,CAAC;IAGM,KAAK;QACR,IAAI,CAAC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,EAAE;YACxB,KAAK,CAAC,KAAK,EAAE,CAAC;QAClB,CAAC,CAAC,CAAA;IACN,CAAC;IAGM,OAAO;QACV,IAAI,CAAC,MAAM,CAAC,OAAO,CAAC,KAAK,CAAC,EAAE;YACxB,KAAK,CAAC,OAAO,EAAE,CAAC;QACpB,CAAC,CAAC,CAAA;IACN,CAAC;IAGD,IAAW,IAAI;QACX,gFAAgF;QAChF,OAAO,IAAI,CAAC,MAAM,CAAC,OAAO,EAAE,CAAC,IAAI,EAAE,CAAC,KAAK,EAAE,CAAC,CAAC,CAAC,CAAC,IAAI,IAAI,CAAC,CAAC;IAC7D,CAAC;IAGD,IAAW,iBAAiB;QACxB,OAAO,IAAI,CAAC,mBAAmB,CAAC;IACpC,CAAC;CACJ;AAGD,MAAM,OAAO,OAAO;IAEN,SAAS,CAA0B;IACnC,WAAW,CAAyB;IAE9C,uFAAuF;IAC7E,gBAAgB,GAAW,CAAC,CAAC;IAE7B,UAAU,CAAS;IACnB,mBAAmB,CAAS;IAC5B,YAAY,CAAS;IACrB,QAAQ,CAAS;IAE3B,YAAY,EAAE,SAAS,EAAE,iBAAiB,EAAE,QAAQ,EAAE,OAAO,EAAE,KAAK,GAAG,SAAS,EAAe;QAC3F,MAAM,WAAW,GAAG,CAAC,SAAS,EAAE,QAAQ,EAAE,iBAAiB,EAAE,OAAO,CAAqC,CAAC;QAE1G,IAAI,CAAC,SAAS,GAAG,EAAE,CAAC,QAAQ,CAAC,EAAE,CAAC,KAAK,CAAC,WAAW,EAAE,KAAK,CAAC,EAAE,KAAK,CAAC,CAAC;QAClE,IAAI,CAAC,WAAW,GAAG,EAAE,CAAC,QAAQ,CAAC,EAAE,CAAC,KAAK,CAAC,WAAW,EAAE,KAAK,CAAC,EAAE,KAAK,CAAC,CAAC;QAEpE,IAAI,CAAC,UAAU,GAAG,SAAS,CAAC;QAC5B,IAAI,CAAC,mBAAmB,GAAG,iBAAiB,CAAC;QAC7C,IAAI,CAAC,YAAY,GAAG,QAAQ,CAAC;QAC7B,IAAI,CAAC,QAAQ,GAAG,OAAO,CAAC;IAC5B,CAAC;IAGD;;OAEG;IACI,MAAM,CAAC,GAAgB,EAAE,KAAkB;QAC9C,MAAM,UAAU,GAAG,GAAG,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;QAChC,MAAM,OAAO,GAAG,GAAG,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;QAE7B,IAAI,UAAU,GAAG,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC;YACvC,MAAM,KAAK,CAAC,2DAA2D;gBACnE,IAAI,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC,CAAC,CAAC,+CAA+C,UAAU,EAAE,CAAC,CAAA;QAC/F,CAAC;QAED,IAAI,IAAI,CAAC,gBAAgB,GAAG,OAAO,GAAG,IAAI,CAAC,mBAAmB,EAAE,CAAC;YAC7D,MAAM,KAAK,CAAC,4DAA4D,IAAI,CAAC,mBAAmB,uBAAuB,CAAC,CAAC;QAC7H,CAAC;QAED,MAAM,aAAa,GAAG,IAAI,CAAC,cAAc,CAAC,GAAG,EAAE,IAAI,CAAC,SAAS,CAAC,CAAC;QAC/D,MAAM,eAAe,GAAG,IAAI,CAAC,cAAc,CAAC,KAAK,EAAE,IAAI,CAAC,WAAW,CAAC,CAAC;QAErE,IAAI,CAAC,SAAS,CAAC,MAAM,CAAC,aAAa,CAAC,CAAC;QACrC,IAAI,CAAC,WAAW,CAAC,MAAM,CAAC,eAAe,CAAC,CAAC;QAEzC,aAAa,CAAC,OAAO,EAAE,CAAC;QACxB,eAAe,CAAC,OAAO,EAAE,CAAC;QAE1B,6DAA6D;QAC7D,IAAI,CAAC,gBAAgB,IAAI,OAAO,CAAC;QAEjC,OAAO;YACH,QAAQ,EAAE,IAAI,CAAC,SAAS;YACxB,UAAU,EAAE,IAAI,CAAC,WAAW;SAC/B,CAAA;IACL,CAAC;IAGS,cAAc,CAAC,SAAsB,EAAE,aAA0B;QACvE,MAAM,OAAO,GAAG,SAAS,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;QAEnC,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAEhB,MAAM,UAAU,GAAG,aAAa,CAAC,KAAK,CAClC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EACZ,CAAC,IAAI,CAAC,UAAU,EAAE,IAAI,CAAC,YAAY,EAAE,IAAI,CAAC,gBAAgB,EAAE,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC;YAEhF,MAAM,MAAM,GAAG,aAAa,CAAC,KAAK,CAC9B,CAAC,CAAC,EAAE,CAAC,EAAE,IAAI,CAAC,gBAAgB,GAAG,OAAO,EAAE,CAAC,CAAC,EAC1C,CAAC,IAAI,CAAC,UAAU,EAAE,IAAI,CAAC,YAAY,EAAE,IAAI,CAAC,mBAAmB,GAAG,IAAI,CAAC,gBAAgB,GAAG,OAAO,EAAE,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC;YAErH,8EAA8E;YAC9E,qFAAqF;YACrF,6DAA6D;YAC7D,OAAO,EAAE,CAAC,MAAM,CAAC,CAAC,UAAU,EAAE,SAAS,EAAE,MAAM,CAAC,EAAE,CAAC,CAAC,CAAC;QACzD,CAAC,CAAC,CAAA;IACN,CAAC;IAGM,KAAK;QACR,IAAI,CAAC,gBAAgB,GAAG,CAAC,CAAC;QAE1B,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YACT,MAAM,eAAe,GAAG,IAAI,CAAC,SAAS,CAAC,KAAK,CAAC;YAC7C,MAAM,iBAAiB,GAAG,IAAI,CAAC,WAAW,CAAC,KAAK,CAAC;YAEjD,IAAI,CAAC,SAAS,CAAC,MAAM,CAAC,EAAE,CAAC,KAAK,CAAC,eAAe,CAAC,CAAC,CAAC;YACjD,IAAI,CAAC,WAAW,CAAC,MAAM,CAAC,EAAE,CAAC,KAAK,CAAC,iBAAiB,CAAC,CAAC,CAAC;QACzD,CAAC,CAAC,CAAC;IACP,CAAC;IAGM,OAAO;QACV,IAAI,CAAC,SAAS,CAAC,OAAO,EAAE,CAAC;QACzB,IAAI,CAAC,WAAW,CAAC,OAAO,EAAE,CAAC;IAC/B,CAAC;IAGD;;OAEG;IACH,IAAI,IAAI;QACJ,OAAO,IAAI,CAAC,gBAAgB,CAAC;IACjC,CAAC;CAEJ"}

package/dist/layers/cached_rope_multihead_attention.d.ts

@@ -0,0 +1,31 @@
+import * as tf from '@tensorflow/tfjs';
+import { KvCacheContainer } from "../kv_cache";
+import { MultiHeadAttention, type MultiHeadAttentionArgs } from '../layers/multihead_attention';
+import { type Kwargs } from '@tensorflow/tfjs-layers/dist/types';
+/**
+ * MultiHeadAttention with RoPE and KV caching. If using KV caching, this layer
+ * should be used in a custom training loop because it requires the cache to be
+ * passed through the `kwargs.kvCache` argument during the `layer.apply()`
+ * forward propagation.
+ *
+ * If a KV cache is not provided, then this layer operates as MultiHeadAttention with RoPE.
+ */
+export declare class CachedRoPEMultiHeadAttention extends MultiHeadAttention {
+    static className: string;
+    protected rope: tf.layers.Layer;
+    constructor(args: MultiHeadAttentionArgs);
+    protected forward(query_input: tf.Tensor, key_input: tf.Tensor, value_input: tf.Tensor, packing_mask: tf.Tensor | null, causal_mask: tf.Tensor | null, kwargs: Kwargs): tf.Tensor;
+    protected getCachedKV(kv_container: KvCacheContainer, key_split: tf.Tensor4D, value_split: tf.Tensor4D): {
+        keyCache: tf.Variable<tf.Rank.R4>;
+        valueCache: tf.Variable<tf.Rank.R4>;
+    };
+    /**
+     * Adds RoPE position encoding right after splitting heads.
+     */
+    protected splitHeads(query: tf.Tensor, key: tf.Tensor, value: tf.Tensor, shuffle: number[]): {
+        query_split: tf.Tensor4D;
+        key_split: tf.Tensor4D;
+        value_split: tf.Tensor4D;
+    };
+}
+//# sourceMappingURL=cached_rope_multihead_attention.d.ts.map