@stellarapp/tfjs-stellar 1.0.0

package/jest.config.ts

@@ -0,0 +1,203 @@
+/**
+ * For a detailed explanation regarding each configuration property, visit:
+ * https://jestjs.io/docs/configuration
+ */
+
+import type { Config } from 'jest';
+
+const config: Config = {
+    setupFiles: [],
+
+    extensionsToTreatAsEsm: [".ts"],
+
+    // A map from regular expressions to paths to transformers
+    transform: {
+        "^.+\.ts?$": ["ts-jest", {
+            useESM: true
+        }],
+    },
+
+    // An array of regexp pattern strings that are matched against all test paths, matched tests are skipped
+    testPathIgnorePatterns: [
+        "/node_modules/", "e2e"
+    ],
+
+    // 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",
+    //   "text",
+    //   "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",
+    //   "mjs",
+    //   "cjs",
+    //   "jsx",
+    //   "ts",
+    //   "tsx",
+    //   "json",
+    //   "node"
+    // ],
+
+    // An array of regexp pattern strings, matched against all module paths before considered 'visible' to the module loader
+    // modulePathIgnorePatterns: [],
+
+    // 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;

package/package.json

@@ -0,0 +1,24 @@
+{
+  "name": "@stellarapp/tfjs-stellar",
+  "version": "1.0.0",
+  "description": "An extension of TensorFlow.js for implementing large language models.",
+  "license": "ISC",
+  "author": "",
+  "type": "module",
+  "main": "index.ts",
+  "scripts": {
+    "test": "npx jest"
+  },
+  "devDependencies": {
+    "@tensorflow/tfjs": "^4.22.0",
+    "@types/jest": "^30.0.0",
+    "@types/node": "^26.0.0",
+    "jest": "^30.4.2",
+    "ts-jest": "^29.4.11",
+    "tsx": "^4.22.4",
+    "typescript": "^6.0.3"
+  },
+  "peerDependencies": {
+    "@tensorflow/tfjs": "*"
+  }
+}

package/src/index.ts

@@ -0,0 +1,93 @@
+export * as models from "./models";
+export * as losses from "./losses";
+export * as metrics from "./metrics";
+
+import { MultiHeadAttention, type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
+export { MultiHeadAttention, type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
+
+import { CachedRoPEMultiHeadAttention } from "@/layers/cached_rope_multihead_attention";
+export { CachedRoPEMultiHeadAttention } from "@/layers/cached_rope_multihead_attention";
+
+import { TransformerEncoder, type TransformerEncoderArgs, } from "@/layers/transformer_encoder";
+export { TransformerEncoder, type TransformerEncoderArgs, } from "@/layers/transformer_encoder";
+
+import { TransformerDecoder, type TransformerDecoderArgs, } from "@/layers/transformer_decoder";
+export { TransformerDecoder, type TransformerDecoderArgs, } from "@/layers/transformer_decoder";
+
+import { TokenAndPositionalEmbedding, type TokenAndPositionalEmbeddingArgs } from "@/layers/token_and_positional_embedding";
+export { TokenAndPositionalEmbedding, type TokenAndPositionalEmbeddingArgs } from "@/layers/token_and_positional_embedding";
+
+import { PositionalEncoding, type PositionalEncodingArgs } from "@/layers/positional_encoding";
+export { PositionalEncoding, type PositionalEncodingArgs } from "@/layers/positional_encoding";
+
+import { GPT2DecoderBlock, type GPTDecoderBlockArgs } from "@/layers/gpt_decoder_block";
+export { GPT2DecoderBlock as GPTDecoderBlock, type GPTDecoderBlockArgs } from "@/layers/gpt_decoder_block";
+
+import { LlmModel, type LlmModelArgs } from "@/models/llm_model";
+export { LlmModel, type LlmModelArgs } from "@/models/llm_model";
+
+import { UNetModel, type UNetModelArgs } from "@/models/u_net";
+
+import { RotaryPositionEmbedding, type RotaryPositionEmbeddingArgs } from "@/layers/rotary_position_embedding";
+export { RotaryPositionEmbedding, type RotaryPositionEmbeddingArgs } from "@/layers/rotary_position_embedding";
+
+
+import { GptModel, type GptModelArgs } from "@/models/gpt_model";
+export { GptModel, type GptModelArgs } from "@/models/gpt_model";
+
+
+// The following exports give a keras-like import just like TFJS's tf.layers.<...>
+
+export function llmModel(args: LlmModelArgs) {
+    return new LlmModel(args);
+}
+
+
+export function gptModel(args: GptModelArgs) {
+    return new GptModel(args);
+}
+
+
+export function tokenAndPositionalEmbedding(args: TokenAndPositionalEmbeddingArgs) {
+    return new TokenAndPositionalEmbedding(args);
+}
+
+
+export function transformerEncoder(args: TransformerEncoderArgs) {
+    return new TransformerEncoder(args);
+}
+
+
+export function transformerDecoder(args: TransformerDecoderArgs) {
+    return new TransformerDecoder(args);
+}
+
+
+export function multiheadAttention(args: MultiHeadAttentionArgs) {
+    return new MultiHeadAttention(args);
+}
+
+
+export function cachedRopeMultiheadAttention(args: MultiHeadAttentionArgs) {
+    return new CachedRoPEMultiHeadAttention(args);
+}
+
+
+export function positionalEncoding(args: PositionalEncodingArgs) {
+    return new PositionalEncoding(args);
+}
+
+
+export function gpt2DecoderBlock(args: GPTDecoderBlockArgs) {
+    return new GPT2DecoderBlock(args);
+}
+
+
+export function unetModel(args: UNetModelArgs) {
+    return new UNetModel(args);
+}
+
+
+export function rotaryPositionEmbedding(args: RotaryPositionEmbeddingArgs) {
+    return new RotaryPositionEmbedding(args);
+}

package/src/kv_cache.ts

@@ -0,0 +1,205 @@
+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) {
+        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">) {
+        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) {
+        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]]);
+
+            return {
+                keyCache: k_cache,
+                valueCache: v_cache
+            }
+        })
+    }
+
+
+    public reset() {
+        this.caches.forEach(cache => {
+            cache.reset();
+        })
+    }
+
+
+    public dispose() {
+        this.caches.forEach(cache => {
+            cache.dispose();
+        })
+    }
+
+
+    public 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() {
+        return this.max_sequence_length;
+    }
+}
+
+
+export class KvCache {
+
+    protected key_cache: tf.Variable<tf.Rank.R4>;
+    protected value_cache: tf.Variable<tf.Rank.R4>
+
+    // 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];
+
+        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) {
+        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}`)
+        }
+
+        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) {
+        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]);
+
+            // 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 {
+        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 {
+        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 {
+        return this.current_position;
+    }
+
+}

package/src/layers/cached_rope_multihead_attention.test.ts

@@ -0,0 +1,59 @@
+import * as tf from '@tensorflow/tfjs';
+
+import { KvCacheContainer } from '@/kv_cache';
+import { CachedRoPEMultiHeadAttention } from '@/layers/cached_rope_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("CachedRoPEMultiHeadAttention tests", () => {
+
+    test("aggregate forward passes output are identical normal multihead attention", () => {
+        compareNormalWithCachedAttention(tf.randomUniform<tf.Rank.R3>([2, 10, 16]), 123);
+        compareNormalWithCachedAttention(tf.randomUniform<tf.Rank.R3>([1, 10, 16]), 123);
+        compareNormalWithCachedAttention(tf.randomUniform<tf.Rank.R3>([1, 1, 16]), 123);
+        compareNormalWithCachedAttention(tf.randomUniform<tf.Rank.R3>([3, 2, 16]), 123);
+
+        // input exceeds KV cach size
+        expect(() => compareNormalWithCachedAttention(tf.randomUniform<tf.Rank.R3>([1, 10, 16]), 5)).toThrow();
+
+        function compareNormalWithCachedAttention(input: tf.Tensor3D, max_sequence_length: number) {
+            const embed_dim = input.shape[2];
+            const batch = input.shape[0];
+            const heads = 2;
+
+            const kv_cache = new KvCacheContainer(max_sequence_length);
+
+            const normal_mha = new CachedRoPEMultiHeadAttention({ numHeads: heads, embedDim: embed_dim, causal: true });
+            const normal_mha_output = normal_mha.apply(input) as tf.Tensor;
+
+            // initialize cached attention with identical configuration and weights
+            const cached_mha1 = new CachedRoPEMultiHeadAttention({ ...normal_mha.getConfig(), name: "cache_test1" });
+            cached_mha1.build(input.shape);
+            cached_mha1.setWeights(normal_mha.getWeights());
+
+            const cached_mha2 = new CachedRoPEMultiHeadAttention({ ...normal_mha.getConfig(), name: "cache_test2" });
+            cached_mha2.build(input.shape);
+            cached_mha2.setWeights(normal_mha.getWeights());
+
+            const cached_mha_outputs1: tf.Tensor[] = [];
+            const cached_mha_outputs2: tf.Tensor[] = [];
+
+            for (let i = 0; i < input.shape[1]; i++) {
+                const current_token = input.slice([0, i, 0], [batch, 1, embed_dim]);
+
+                cached_mha_outputs1.push(cached_mha1.apply(current_token, { kvCache: kv_cache }) as tf.Tensor);
+                cached_mha_outputs2.push(cached_mha2.apply(current_token, { kvCache: kv_cache }) as tf.Tensor);
+            }
+
+            expect(kv_cache.size== input.shape[1]);
+            expect(kv_cache.size == input.shape[1]);
+
+            expect(normal_mha_output.sub(tf.concat(cached_mha_outputs1, 1)).sum().dataSync()[0]).toBeLessThan(1e-6);
+            expect(normal_mha_output.sub(tf.concat(cached_mha_outputs2, 1)).sum().dataSync()[0]).toBeLessThan(1e-6);
+        }
+    })
+});