@stellarapp/tfjs-stellar 1.0.4 → 1.0.6

package/dist/src/kv_cache.js

@@ -1,135 +0,0 @@
-import * as tf from "@tensorflow/tfjs";
-/**
- * A container for KV caches. A model should initialize one KV cache
- */
-export class KvCacheContainer {
-    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;
-    }
-    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]`
-     */
-    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]]);
-            return {
-                keyCache: k_cache,
-                valueCache: v_cache
-            };
-        });
-    }
-    reset() {
-        this.caches.forEach(cache => {
-            cache.reset();
-        });
-    }
-    dispose() {
-        this.caches.forEach(cache => {
-            cache.dispose();
-        });
-    }
-    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;
-    }
-    get maxSequenceLength() {
-        return this.max_sequence_length;
-    }
-}
-export class KvCache {
-    key_cache;
-    value_cache;
-    // the size of the KV cache, represents the number of tokens since the first chat token
-    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]`
-     */
-    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}`);
-        }
-        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,
-        };
-    }
-    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]);
-            // 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);
-        });
-    }
-    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));
-        });
-    }
-    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() {
-        return this.current_position;
-    }
-}
-//# sourceMappingURL=kv_cache.js.map

package/dist/src/kv_cache.js.map

@@ -1 +0,0 @@
-{"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/src/layers/cached_rope_multihead_attention.d.ts

@@ -1,31 +0,0 @@
-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

package/dist/src/layers/cached_rope_multihead_attention.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"cached_rope_multihead_attention.d.ts","sourceRoot":"","sources":["../../../src/layers/cached_rope_multihead_attention.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAC;AAC9C,OAAO,EAAE,kBAAkB,EAAE,KAAK,sBAAsB,EAAE,MAAM,8BAA8B,CAAC;AAE/F,OAAO,EAAE,KAAK,MAAM,EAAE,MAAM,oCAAoC,CAAC;AAGjE;;;;;;;GAOG;AACH,qBAAa,4BAA6B,SAAQ,kBAAkB;IAChE,MAAM,CAAC,SAAS,SAAkC;IAElD,SAAS,CAAC,IAAI,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;gBAEpB,IAAI,EAAE,sBAAsB;cAMrB,OAAO,CACtB,WAAW,EAAE,EAAE,CAAC,MAAM,EACtB,SAAS,EAAE,EAAE,CAAC,MAAM,EACpB,WAAW,EAAE,EAAE,CAAC,MAAM,EACtB,YAAY,EAAE,EAAE,CAAC,MAAM,GAAG,IAAI,EAC9B,WAAW,EAAE,EAAE,CAAC,MAAM,GAAG,IAAI,EAC7B,MAAM,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM;IAuC9B,SAAS,CAAC,WAAW,CAAC,YAAY,EAAE,gBAAgB,EAAE,SAAS,EAAE,EAAE,CAAC,QAAQ,EAAE,WAAW,EAAE,EAAE,CAAC,QAAQ;;;;IAqBtG;;OAEG;cACgB,UAAU,CAAC,KAAK,EAAE,EAAE,CAAC,MAAM,EAAE,GAAG,EAAE,EAAE,CAAC,MAAM,EAAE,KAAK,EAAE,EAAE,CAAC,MAAM,EAAE,OAAO,EAAE,MAAM,EAAE;qBAO5D,EAAE,CAAC,QAAQ;mBAEX,EAAE,CAAC,QAAQ;qBACwB,EAAE,CAAC,QAAQ;;CAIxF"}

package/dist/src/layers/cached_rope_multihead_attention.js

@@ -1,76 +0,0 @@
-import * as tf from '@tensorflow/tfjs';
-import { MultiHeadAttention } from '@/layers/multihead_attention';
-import { RotaryPositionEmbedding } from '@/layers/rotary_position_embedding';
-/**
- * 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 class CachedRoPEMultiHeadAttention extends MultiHeadAttention {
-    static className = "CachedRoPEMultiHeadAttention";
-    rope;
-    constructor(args) {
-        super(args);
-        this.rope = new RotaryPositionEmbedding({ dim: Math.floor(this.embedDim / this.numHeads) });
-    }
-    forward(query_input, key_input, value_input, packing_mask, causal_mask, kwargs) {
-        return tf.tidy(() => {
-            const { query, key, value } = this.applyInputProjections(query_input, key_input, value_input);
-            // swap the seq and heads dimensions: [batch, seq, heads, head_dim] -> [batch, heads, seq, head_dim]
-            const move_head_dim_forward = [0, 2, 1, 3];
-            const split = this.splitHeads(query, key, value, move_head_dim_forward);
-            const query_split = split.query_split;
-            let key_split = split.key_split;
-            let value_split = split.value_split;
-            if (kwargs.training !== true && kwargs.kvCache) {
-                // runs on inference, updates the KV cache and get the historical key and value
-                const cached_kv = this.getCachedKV(kwargs.kvCache, key_split, value_split);
-                key_split = cached_kv.keyCache;
-                value_split = cached_kv.valueCache;
-            }
-            // apply scaled dot production attention to get [batch, seq, numHeads, embedDim]
-            const spda = MultiHeadAttention.scaledDotProductionAttention(query_split, key_split, value_split, kwargs.attentionMask ?? null, packing_mask, causal_mask, this.dropout, this.causal, kwargs);
-            // concat heads and apply the output projection
-            const output = this.outputProjection.apply(spda.transpose(move_head_dim_forward).reshape([query_input.shape[0], query_input.shape[1], this.embedDim]));
-            return output;
-        });
-    }
-    getCachedKV(kv_container, key_split, value_split) {
-        try {
-            let kv_cache = kv_container.update(this.name, key_split, value_split);
-            if (!kv_cache) {
-                kv_container.create(this.name, {
-                    batchSize: key_split.shape[0],
-                    numHeads: this.numHeads,
-                    headDim: this.embedDim / this.numHeads,
-                });
-                kv_cache = kv_container.update(this.name, key_split, value_split);
-            }
-            return kv_cache;
-        }
-        catch (error) {
-            throw Error(`${this.getClassName()}::getCachedKV ${this.name} ${error.toString()}`);
-        }
-    }
-    /**
-     * Adds RoPE position encoding right after splitting heads.
-     */
-    splitHeads(query, key, value, shuffle) {
-        const batch_size = query.shape[0];
-        const split_heads = [batch_size, -1, this.numHeads, this.embedDim / this.numHeads];
-        return tf.tidy(() => {
-            return {
-                query_split: this.rope.apply(query.reshape(split_heads))
-                    .transpose(shuffle),
-                key_split: this.rope.apply(key.reshape(split_heads))
-                    .transpose(shuffle),
-                value_split: value.reshape(split_heads).transpose(shuffle)
-            };
-        });
-    }
-}
-tf.serialization.registerClass(CachedRoPEMultiHeadAttention);
-//# sourceMappingURL=cached_rope_multihead_attention.js.map

package/dist/src/layers/cached_rope_multihead_attention.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"cached_rope_multihead_attention.js","sourceRoot":"","sources":["../../../src/layers/cached_rope_multihead_attention.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAEvC,OAAO,EAAE,kBAAkB,EAA+B,MAAM,8BAA8B,CAAC;AAC/F,OAAO,EAAE,uBAAuB,EAAE,MAAM,oCAAoC,CAAC;AAI7E;;;;;;;GAOG;AACH,MAAM,OAAO,4BAA6B,SAAQ,kBAAkB;IAChE,MAAM,CAAC,SAAS,GAAG,8BAA8B,CAAC;IAExC,IAAI,CAAkB;IAEhC,YAAY,IAA4B;QACpC,KAAK,CAAC,IAAI,CAAC,CAAC;QACZ,IAAI,CAAC,IAAI,GAAG,IAAI,uBAAuB,CAAC,EAAE,GAAG,EAAE,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,QAAQ,CAAC,EAAE,CAAC,CAAC;IAChG,CAAC;IAGkB,OAAO,CACtB,WAAsB,EACtB,SAAoB,EACpB,WAAsB,EACtB,YAA8B,EAC9B,WAA6B,EAC7B,MAAc;QAEd,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,EAAE,KAAK,EAAE,GAAG,EAAE,KAAK,EAAE,GAAG,IAAI,CAAC,qBAAqB,CAAC,WAAW,EAAE,SAAS,EAAE,WAAW,CAAC,CAAC;YAE9F,oGAAoG;YACpG,MAAM,qBAAqB,GAAG,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC;YAE3C,MAAM,KAAK,GAAG,IAAI,CAAC,UAAU,CAAC,KAAK,EAAE,GAAG,EAAE,KAAK,EAAE,qBAAqB,CAAC,CAAC;YAExE,MAAM,WAAW,GAAG,KAAK,CAAC,WAAW,CAAC;YACtC,IAAI,SAAS,GAAG,KAAK,CAAC,SAAS,CAAC;YAChC,IAAI,WAAW,GAAG,KAAK,CAAC,WAAW,CAAC;YAEpC,IAAI,MAAM,CAAC,QAAQ,KAAK,IAAI,IAAI,MAAM,CAAC,OAAO,EAAE,CAAC;gBAC7C,+EAA+E;gBAC/E,MAAM,SAAS,GAAG,IAAI,CAAC,WAAW,CAC9B,MAAM,CAAC,OAA2B,EAAE,SAAS,EAAE,WAAW,CAAC,CAAC;gBAEhE,SAAS,GAAG,SAAS,CAAC,QAAQ,CAAC;gBAC/B,WAAW,GAAG,SAAS,CAAC,UAAU,CAAC;YACvC,CAAC;YAED,gFAAgF;YAChF,MAAM,IAAI,GAAG,kBAAkB,CAAC,4BAA4B,CACxD,WAAW,EAAE,SAAS,EAAE,WAAW,EACnC,MAAM,CAAC,aAAa,IAAI,IAAI,EAAE,YAAY,EAAE,WAAW,EACvD,IAAI,CAAC,OAAO,EAAE,IAAI,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;YAEvC,+CAA+C;YAC/C,MAAM,MAAM,GAAG,IAAI,CAAC,gBAAgB,CAAC,KAAK,CACtC,IAAI,CAAC,SAAS,CAAC,qBAAqB,CAAC,CAAC,OAAO,CACzC,CAAC,WAAW,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,WAAW,CAAC,KAAK,CAAC,CAAC,CAAE,EAAE,IAAI,CAAC,QAAQ,CAAC,CAAC,CAAC,CAAC;YAEvE,OAAO,MAAmB,CAAC;QAC/B,CAAC,CAAC,CAAA;IACN,CAAC;IAGS,WAAW,CAAC,YAA8B,EAAE,SAAsB,EAAE,WAAwB;QAClG,IAAI,CAAC;YACD,IAAI,QAAQ,GAAG,YAAY,CAAC,MAAM,CAAC,IAAI,CAAC,IAAI,EAAE,SAAS,EAAE,WAAW,CAAC,CAAC;YAEtE,IAAI,CAAC,QAAQ,EAAE,CAAC;gBACZ,YAAY,CAAC,MAAM,CAAC,IAAI,CAAC,IAAI,EAAE;oBAC3B,SAAS,EAAE,SAAS,CAAC,KAAK,CAAC,CAAC,CAAC;oBAC7B,QAAQ,EAAE,IAAI,CAAC,QAAQ;oBACvB,OAAO,EAAE,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,QAAQ;iBACzC,CAAC,CAAA;gBAEF,QAAQ,GAAG,YAAY,CAAC,MAAM,CAAC,IAAI,CAAC,IAAI,EAAE,SAAS,EAAE,WAAW,CAAE,CAAC;YACvE,CAAC;YAED,OAAO,QAAS,CAAC;QACrB,CAAC;QAAC,OAAO,KAAU,EAAE,CAAC;YAClB,MAAM,KAAK,CAAC,GAAG,IAAI,CAAC,YAAY,EAAE,iBAAiB,IAAI,CAAC,IAAI,IAAI,KAAK,CAAC,QAAQ,EAAE,EAAE,CAAC,CAAC;QACxF,CAAC;IACL,CAAC;IAGD;;OAEG;IACgB,UAAU,CAAC,KAAgB,EAAE,GAAc,EAAE,KAAgB,EAAE,OAAiB;QAC/F,MAAM,UAAU,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;QAClC,MAAM,WAAW,GAAG,CAAC,UAAU,EAAE,CAAC,CAAC,EAAE,IAAI,CAAC,QAAQ,EAAE,IAAI,CAAC,QAAQ,GAAG,IAAI,CAAC,QAAQ,CAAC,CAAC;QAEnF,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,OAAO;gBACH,WAAW,EAAG,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,CAAe;qBAClE,SAAS,CAAC,OAAO,CAAgB;gBACtC,SAAS,EAAG,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,GAAG,CAAC,OAAO,CAAC,WAAW,CAAC,CAAe;qBAC9D,SAAS,CAAC,OAAO,CAAgB;gBACtC,WAAW,EAAE,KAAK,CAAC,OAAO,CAAC,WAAW,CAAC,CAAC,SAAS,CAAC,OAAO,CAAgB;aAC5E,CAAA;QACL,CAAC,CAAC,CAAA;IACN,CAAC;;AAIL,EAAE,CAAC,aAAa,CAAC,aAAa,CAAC,4BAA4B,CAAC,CAAC"}

package/dist/src/layers/cached_rope_multihead_attention.test.d.ts

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

package/dist/src/layers/cached_rope_multihead_attention.test.d.ts.map

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

package/dist/src/layers/cached_rope_multihead_attention.test.js

@@ -1,43 +0,0 @@
-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([2, 10, 16]), 123);
-        compareNormalWithCachedAttention(tf.randomUniform([1, 10, 16]), 123);
-        compareNormalWithCachedAttention(tf.randomUniform([1, 1, 16]), 123);
-        compareNormalWithCachedAttention(tf.randomUniform([3, 2, 16]), 123);
-        // input exceeds KV cach size
-        expect(() => compareNormalWithCachedAttention(tf.randomUniform([1, 10, 16]), 5)).toThrow();
-        function compareNormalWithCachedAttention(input, max_sequence_length) {
-            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);
-            // 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 = [];
-            const cached_mha_outputs2 = [];
-            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 }));
-                cached_mha_outputs2.push(cached_mha2.apply(current_token, { kvCache: kv_cache }));
-            }
-            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);
-        }
-    });
-});
-//# sourceMappingURL=cached_rope_multihead_attention.test.js.map

package/dist/src/layers/cached_rope_multihead_attention.test.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"cached_rope_multihead_attention.test.js","sourceRoot":"","sources":["../../../src/layers/cached_rope_multihead_attention.test.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAEvC,OAAO,EAAE,gBAAgB,EAAE,MAAM,YAAY,CAAC;AAC9C,OAAO,EAAE,4BAA4B,EAAE,MAAM,0CAA0C,CAAC;AAGxF,sDAAsD;AACtD,wEAAwE;AACxE,EAAE,CAAC,GAAG,EAAE,CAAC,GAAG,CAAC,SAAS,EAAE,KAAK,CAAC,CAAC;AAG/B,QAAQ,CAAC,oCAAoC,EAAE,GAAG,EAAE;IAEhD,IAAI,CAAC,0EAA0E,EAAE,GAAG,EAAE;QAClF,gCAAgC,CAAC,EAAE,CAAC,aAAa,CAAa,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;QACjF,gCAAgC,CAAC,EAAE,CAAC,aAAa,CAAa,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;QACjF,gCAAgC,CAAC,EAAE,CAAC,aAAa,CAAa,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;QAChF,gCAAgC,CAAC,EAAE,CAAC,aAAa,CAAa,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,EAAE,GAAG,CAAC,CAAC;QAEhF,6BAA6B;QAC7B,MAAM,CAAC,GAAG,EAAE,CAAC,gCAAgC,CAAC,EAAE,CAAC,aAAa,CAAa,CAAC,CAAC,EAAE,EAAE,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC,CAAC,OAAO,EAAE,CAAC;QAEvG,SAAS,gCAAgC,CAAC,KAAkB,EAAE,mBAA2B;YACrF,MAAM,SAAS,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;YACjC,MAAM,KAAK,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;YAC7B,MAAM,KAAK,GAAG,CAAC,CAAC;YAEhB,MAAM,QAAQ,GAAG,IAAI,gBAAgB,CAAC,mBAAmB,CAAC,CAAC;YAE3D,MAAM,UAAU,GAAG,IAAI,4BAA4B,CAAC,EAAE,QAAQ,EAAE,KAAK,EAAE,QAAQ,EAAE,SAAS,EAAE,MAAM,EAAE,IAAI,EAAE,CAAC,CAAC;YAC5G,MAAM,iBAAiB,GAAG,UAAU,CAAC,KAAK,CAAC,KAAK,CAAc,CAAC;YAE/D,uEAAuE;YACvE,MAAM,WAAW,GAAG,IAAI,4BAA4B,CAAC,EAAE,GAAG,UAAU,CAAC,SAAS,EAAE,EAAE,IAAI,EAAE,aAAa,EAAE,CAAC,CAAC;YACzG,WAAW,CAAC,KAAK,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC;YAC/B,WAAW,CAAC,UAAU,CAAC,UAAU,CAAC,UAAU,EAAE,CAAC,CAAC;YAEhD,MAAM,WAAW,GAAG,IAAI,4BAA4B,CAAC,EAAE,GAAG,UAAU,CAAC,SAAS,EAAE,EAAE,IAAI,EAAE,aAAa,EAAE,CAAC,CAAC;YACzG,WAAW,CAAC,KAAK,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC;YAC/B,WAAW,CAAC,UAAU,CAAC,UAAU,CAAC,UAAU,EAAE,CAAC,CAAC;YAEhD,MAAM,mBAAmB,GAAgB,EAAE,CAAC;YAC5C,MAAM,mBAAmB,GAAgB,EAAE,CAAC;YAE5C,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;gBACtC,MAAM,aAAa,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,CAAC,CAAC,EAAE,CAAC,KAAK,EAAE,CAAC,EAAE,SAAS,CAAC,CAAC,CAAC;gBAEpE,mBAAmB,CAAC,IAAI,CAAC,WAAW,CAAC,KAAK,CAAC,aAAa,EAAE,EAAE,OAAO,EAAE,QAAQ,EAAE,CAAc,CAAC,CAAC;gBAC/F,mBAAmB,CAAC,IAAI,CAAC,WAAW,CAAC,KAAK,CAAC,aAAa,EAAE,EAAE,OAAO,EAAE,QAAQ,EAAE,CAAc,CAAC,CAAC;YACnG,CAAC;YAED,MAAM,CAAC,QAAQ,CAAC,IAAI,IAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC;YACvC,MAAM,CAAC,QAAQ,CAAC,IAAI,IAAI,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC,CAAC;YAExC,MAAM,CAAC,iBAAiB,CAAC,GAAG,CAAC,EAAE,CAAC,MAAM,CAAC,mBAAmB,EAAE,CAAC,CAAC,CAAC,CAAC,GAAG,EAAE,CAAC,QAAQ,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,YAAY,CAAC,IAAI,CAAC,CAAC;YACxG,MAAM,CAAC,iBAAiB,CAAC,GAAG,CAAC,EAAE,CAAC,MAAM,CAAC,mBAAmB,EAAE,CAAC,CAAC,CAAC,CAAC,GAAG,EAAE,CAAC,QAAQ,EAAE,CAAC,CAAC,CAAC,CAAC,CAAC,YAAY,CAAC,IAAI,CAAC,CAAC;QAC5G,CAAC;IACL,CAAC,CAAC,CAAA;AACN,CAAC,CAAC,CAAC"}

package/dist/src/layers/gpt_decoder_block.d.ts

@@ -1,34 +0,0 @@
-import * as tf from "@tensorflow/tfjs";
-import { type Kwargs } from "@tensorflow/tfjs-layers/dist/types";
-import { type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
-import { TransformerDecoder, type TransformerDecoderArgs } from "@/layers/transformer_decoder";
-export interface GPTDecoderBlockArgs extends Omit<MultiHeadAttentionArgs, "causal"> {
-    dimsFeedForward?: number;
-}
-/**
- * This implements the GPT-2 transformer block by modifying the transformer
- * decoder block to use pre-layer-normalization and replacing ReLU activation
- * with GELU.
- *
- * @param numHeads number of attention heads to use
- * @param embedDim the embedding size of the input (input embeddings, typically the last dimension)
- * @param causal use causal masking on inputs (masks future inputs to prevent looking ahead), default `true`
- * @param dropout use dropout during the attention calculations, default `0.1`
- * @param dimsFeedForward the size of the intermediate feed forward layer, default `2048`
- * @param useBias use bias for the dense sublayers and multiHead attention's dense sublayers, default `true`
- */
-export declare class GPT2DecoderBlock extends TransformerDecoder {
-    static className: string;
-    constructor(args: TransformerDecoderArgs);
-    /**
-     * Attention sub-block which is similar to the original transformer except
-     * layer normalization is applied beginning
-     */
-    protected causalSelfAttentionBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor;
-    /**
-     * Feedforward sub-block which is similar to the original transformer except
-     * layer normalization is applied at the beginning and gelu activation is used
-     */
-    protected feedForwardBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor;
-}
-//# sourceMappingURL=gpt_decoder_block.d.ts.map

package/dist/src/layers/gpt_decoder_block.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"gpt_decoder_block.d.ts","sourceRoot":"","sources":["../../../src/layers/gpt_decoder_block.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AACvC,OAAO,EAAE,KAAK,MAAM,EAAE,MAAM,oCAAoC,CAAC;AAEjE,OAAO,EAAE,KAAK,sBAAsB,EAAE,MAAM,8BAA8B,CAAC;AAC3E,OAAO,EAAE,kBAAkB,EAAE,KAAK,sBAAsB,EAAE,MAAM,8BAA8B,CAAC;AAG/F,MAAM,WAAW,mBAAoB,SAAQ,IAAI,CAAC,sBAAsB,EAAE,QAAQ,CAAC;IAC/E,eAAe,CAAC,EAAE,MAAM,CAAC;CAC5B;AAGD;;;;;;;;;;;GAWG;AACH,qBAAa,gBAAiB,SAAQ,kBAAkB;IACpD,MAAM,CAAC,SAAS,SAAsB;gBAG1B,IAAI,EAAE,sBAAsB;IAKxC;;;OAGG;cACgB,wBAAwB,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM;IAcpF;;;OAGG;cACgB,gBAAgB,CAAC,CAAC,EAAE,EAAE,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM;CAkB/E"}

package/dist/src/layers/gpt_decoder_block.js

@@ -1,51 +0,0 @@
-import * as tf from "@tensorflow/tfjs";
-import { TransformerDecoder } from "@/layers/transformer_decoder";
-/**
- * This implements the GPT-2 transformer block by modifying the transformer
- * decoder block to use pre-layer-normalization and replacing ReLU activation
- * with GELU.
- *
- * @param numHeads number of attention heads to use
- * @param embedDim the embedding size of the input (input embeddings, typically the last dimension)
- * @param causal use causal masking on inputs (masks future inputs to prevent looking ahead), default `true`
- * @param dropout use dropout during the attention calculations, default `0.1`
- * @param dimsFeedForward the size of the intermediate feed forward layer, default `2048`
- * @param useBias use bias for the dense sublayers and multiHead attention's dense sublayers, default `true`
- */
-export class GPT2DecoderBlock extends TransformerDecoder {
-    static className = "GPT2DecoderBlock";
-    constructor(args) {
-        super(args);
-    }
-    /**
-     * Attention sub-block which is similar to the original transformer except
-     * layer normalization is applied beginning
-     */
-    causalSelfAttentionBlock(x, kwargs) {
-        return tf.tidy(() => {
-            const residual = x;
-            let attention = this.causalSelfAttentionNorm.apply(x, kwargs);
-            attention = this.causalSelfAttention.apply(attention, kwargs);
-            attention = this.causalSelfAttentionDropout.apply(attention, kwargs);
-            attention = tf.add(attention, residual);
-            return attention;
-        });
-    }
-    /**
-     * Feedforward sub-block which is similar to the original transformer except
-     * layer normalization is applied at the beginning and gelu activation is used
-     */
-    feedForwardBlock(x, kwargs) {
-        return tf.tidy(() => {
-            const residual = x;
-            let feedForward = this.feedFowardNorm.apply(x, kwargs);
-            feedForward = this.feedforward1.apply(feedForward, kwargs);
-            feedForward = this.feedforward2.apply(feedForward, kwargs);
-            feedForward = this.feedForwardDropout.apply(feedForward, kwargs);
-            feedForward = tf.add(feedForward, residual);
-            return feedForward;
-        });
-    }
-}
-tf.serialization.registerClass(GPT2DecoderBlock);
-//# sourceMappingURL=gpt_decoder_block.js.map

package/dist/src/layers/gpt_decoder_block.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"gpt_decoder_block.js","sourceRoot":"","sources":["../../../src/layers/gpt_decoder_block.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAIvC,OAAO,EAAE,kBAAkB,EAA+B,MAAM,8BAA8B,CAAC;AAQ/F;;;;;;;;;;;GAWG;AACH,MAAM,OAAO,gBAAiB,SAAQ,kBAAkB;IACpD,MAAM,CAAC,SAAS,GAAG,kBAAkB,CAAC;IAGtC,YAAY,IAA4B;QACpC,KAAK,CAAC,IAAI,CAAC,CAAC;IAChB,CAAC;IAGD;;;OAGG;IACgB,wBAAwB,CAAC,CAAY,EAAE,MAAc;QACpE,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,QAAQ,GAAG,CAAC,CAAC;YAEnB,IAAI,SAAS,GAAG,IAAI,CAAC,uBAAuB,CAAC,KAAK,CAAC,CAAC,EAAE,MAAM,CAAc,CAAC;YAC3E,SAAS,GAAG,IAAI,CAAC,mBAAmB,CAAC,KAAK,CAAC,SAAS,EAAE,MAAM,CAAc,CAAC;YAC3E,SAAS,GAAG,IAAI,CAAC,0BAA0B,CAAC,KAAK,CAAC,SAAS,EAAE,MAAM,CAAc,CAAC;YAClF,SAAS,GAAG,EAAE,CAAC,GAAG,CAAC,SAAS,EAAE,QAAQ,CAAC,CAAC;YAExC,OAAO,SAAS,CAAC;QACrB,CAAC,CAAC,CAAC;IACP,CAAC;IAGD;;;OAGG;IACgB,gBAAgB,CAAC,CAAY,EAAE,MAAc;QAC5D,OAAO,EAAE,CAAC,IAAI,CAAC,GAAG,EAAE;YAChB,MAAM,QAAQ,GAAG,CAAC,CAAC;YAEnB,IAAI,WAAW,GAAG,IAAI,CAAC,cAAc,CAAC,KAAK,CAAC,CAAC,EAAE,MAAM,CAAC,CAAC;YACvD,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,KAAK,CAAC,WAAW,EAAE,MAAM,CAAC,CAAC;YAC3D,WAAW,GAAG,IAAI,CAAC,YAAY,CAAC,KAAK,CAAC,WAAW,EAAE,MAAM,CAAC,CAAC;YAC3D,WAAW,GAAG,IAAI,CAAC,kBAAkB,CAAC,KAAK,CAAC,WAAW,EAAE,MAAM,CAAc,CAAC;YAC9E,WAAW,GAAG,EAAE,CAAC,GAAG,CAAC,WAAW,EAAE,QAAQ,CAAC,CAAC;YAE5C,OAAO,WAAW,CAAC;QACvB,CAAC,CAAC,CAAC;IACP,CAAC;;AASL,EAAE,CAAC,aAAa,CAAC,aAAa,CAAC,gBAAgB,CAAC,CAAC"}

package/dist/src/layers/index.d.ts

@@ -1,17 +0,0 @@
-import { CachedRoPEMultiHeadAttention } from "./cached_rope_multihead_attention";
-import { GPT2DecoderBlock, GPTDecoderBlockArgs } from "./gpt_decoder_block";
-import { MultiHeadAttention, MultiHeadAttentionArgs } from "./multihead_attention";
-import { PositionalEncoding, PositionalEncodingArgs } from "./positional_encoding";
-import { RotaryPositionEmbedding, RotaryPositionEmbeddingArgs } from "./rotary_position_embedding";
-import { TokenAndPositionalEmbedding, TokenAndPositionalEmbeddingArgs } from "./token_and_positional_embedding";
-import { TransformerDecoder, TransformerDecoderArgs } from "./transformer_decoder";
-import { TransformerEncoder, TransformerEncoderArgs } from "./transformer_encoder";
-export declare function tokenAndPositionalEmbedding(args: TokenAndPositionalEmbeddingArgs): TokenAndPositionalEmbedding;
-export declare function transformerEncoder(args: TransformerEncoderArgs): TransformerEncoder;
-export declare function transformerDecoder(args: TransformerDecoderArgs): TransformerDecoder;
-export declare function multiheadAttention(args: MultiHeadAttentionArgs): MultiHeadAttention;
-export declare function cachedRopeMultiheadAttention(args: MultiHeadAttentionArgs): CachedRoPEMultiHeadAttention;
-export declare function positionalEncoding(args: PositionalEncodingArgs): PositionalEncoding;
-export declare function gpt2DecoderBlock(args: GPTDecoderBlockArgs): GPT2DecoderBlock;
-export declare function rotaryPositionEmbedding(args: RotaryPositionEmbeddingArgs): RotaryPositionEmbedding;
-//# sourceMappingURL=index.d.ts.map

package/dist/src/layers/index.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"index.d.ts","sourceRoot":"","sources":["../../../src/layers/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,4BAA4B,EAAE,MAAM,mCAAmC,CAAC;AACjF,OAAO,EAAE,gBAAgB,EAAE,mBAAmB,EAAE,MAAM,qBAAqB,CAAC;AAC5E,OAAO,EAAE,kBAAkB,EAAE,sBAAsB,EAAE,MAAM,uBAAuB,CAAC;AACnF,OAAO,EAAE,kBAAkB,EAAE,sBAAsB,EAAE,MAAM,uBAAuB,CAAC;AACnF,OAAO,EAAE,uBAAuB,EAAE,2BAA2B,EAAE,MAAM,6BAA6B,CAAC;AACnG,OAAO,EAAE,2BAA2B,EAAE,+BAA+B,EAAE,MAAM,kCAAkC,CAAC;AAChH,OAAO,EAAE,kBAAkB,EAAE,sBAAsB,EAAE,MAAM,uBAAuB,CAAC;AACnF,OAAO,EAAE,kBAAkB,EAAE,sBAAsB,EAAE,MAAM,uBAAuB,CAAC;AAGnF,wBAAgB,2BAA2B,CAAC,IAAI,EAAE,+BAA+B,+BAEhF;AAGD,wBAAgB,kBAAkB,CAAC,IAAI,EAAE,sBAAsB,sBAE9D;AAGD,wBAAgB,kBAAkB,CAAC,IAAI,EAAE,sBAAsB,sBAE9D;AAGD,wBAAgB,kBAAkB,CAAC,IAAI,EAAE,sBAAsB,sBAE9D;AAGD,wBAAgB,4BAA4B,CAAC,IAAI,EAAE,sBAAsB,gCAExE;AAGD,wBAAgB,kBAAkB,CAAC,IAAI,EAAE,sBAAsB,sBAE9D;AAGD,wBAAgB,gBAAgB,CAAC,IAAI,EAAE,mBAAmB,oBAEzD;AAGD,wBAAgB,uBAAuB,CAAC,IAAI,EAAE,2BAA2B,2BAExE"}

package/dist/src/layers/index.js

@@ -1,33 +0,0 @@
-import { CachedRoPEMultiHeadAttention } from "./cached_rope_multihead_attention";
-import { GPT2DecoderBlock } from "./gpt_decoder_block";
-import { MultiHeadAttention } from "./multihead_attention";
-import { PositionalEncoding } from "./positional_encoding";
-import { RotaryPositionEmbedding } from "./rotary_position_embedding";
-import { TokenAndPositionalEmbedding } from "./token_and_positional_embedding";
-import { TransformerDecoder } from "./transformer_decoder";
-import { TransformerEncoder } from "./transformer_encoder";
-export function tokenAndPositionalEmbedding(args) {
-    return new TokenAndPositionalEmbedding(args);
-}
-export function transformerEncoder(args) {
-    return new TransformerEncoder(args);
-}
-export function transformerDecoder(args) {
-    return new TransformerDecoder(args);
-}
-export function multiheadAttention(args) {
-    return new MultiHeadAttention(args);
-}
-export function cachedRopeMultiheadAttention(args) {
-    return new CachedRoPEMultiHeadAttention(args);
-}
-export function positionalEncoding(args) {
-    return new PositionalEncoding(args);
-}
-export function gpt2DecoderBlock(args) {
-    return new GPT2DecoderBlock(args);
-}
-export function rotaryPositionEmbedding(args) {
-    return new RotaryPositionEmbedding(args);
-}
-//# sourceMappingURL=index.js.map

package/dist/src/layers/index.js.map

@@ -1 +0,0 @@
-{"version":3,"file":"index.js","sourceRoot":"","sources":["../../../src/layers/index.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,4BAA4B,EAAE,MAAM,mCAAmC,CAAC;AACjF,OAAO,EAAE,gBAAgB,EAAuB,MAAM,qBAAqB,CAAC;AAC5E,OAAO,EAAE,kBAAkB,EAA0B,MAAM,uBAAuB,CAAC;AACnF,OAAO,EAAE,kBAAkB,EAA0B,MAAM,uBAAuB,CAAC;AACnF,OAAO,EAAE,uBAAuB,EAA+B,MAAM,6BAA6B,CAAC;AACnG,OAAO,EAAE,2BAA2B,EAAmC,MAAM,kCAAkC,CAAC;AAChH,OAAO,EAAE,kBAAkB,EAA0B,MAAM,uBAAuB,CAAC;AACnF,OAAO,EAAE,kBAAkB,EAA0B,MAAM,uBAAuB,CAAC;AAGnF,MAAM,UAAU,2BAA2B,CAAC,IAAqC;IAC7E,OAAO,IAAI,2BAA2B,CAAC,IAAI,CAAC,CAAC;AACjD,CAAC;AAGD,MAAM,UAAU,kBAAkB,CAAC,IAA4B;IAC3D,OAAO,IAAI,kBAAkB,CAAC,IAAI,CAAC,CAAC;AACxC,CAAC;AAGD,MAAM,UAAU,kBAAkB,CAAC,IAA4B;IAC3D,OAAO,IAAI,kBAAkB,CAAC,IAAI,CAAC,CAAC;AACxC,CAAC;AAGD,MAAM,UAAU,kBAAkB,CAAC,IAA4B;IAC3D,OAAO,IAAI,kBAAkB,CAAC,IAAI,CAAC,CAAC;AACxC,CAAC;AAGD,MAAM,UAAU,4BAA4B,CAAC,IAA4B;IACrE,OAAO,IAAI,4BAA4B,CAAC,IAAI,CAAC,CAAC;AAClD,CAAC;AAGD,MAAM,UAAU,kBAAkB,CAAC,IAA4B;IAC3D,OAAO,IAAI,kBAAkB,CAAC,IAAI,CAAC,CAAC;AACxC,CAAC;AAGD,MAAM,UAAU,gBAAgB,CAAC,IAAyB;IACtD,OAAO,IAAI,gBAAgB,CAAC,IAAI,CAAC,CAAC;AACtC,CAAC;AAGD,MAAM,UAAU,uBAAuB,CAAC,IAAiC;IACrE,OAAO,IAAI,uBAAuB,CAAC,IAAI,CAAC,CAAC;AAC7C,CAAC"}

package/dist/src/layers/multihead_attention.d.ts

@@ -1,106 +0,0 @@
-import * as tf from '@tensorflow/tfjs';
-import { type LayerArgs } from '@tensorflow/tfjs-layers/dist/engine/topology';
-import { type Kwargs } from '@tensorflow/tfjs-layers/dist/types';
-export interface MultiHeadAttentionArgs extends LayerArgs {
-    numHeads: number;
-    embedDim: number;
-    useBias?: boolean;
-    dropout?: number;
-    causal?: boolean;
-}
-export interface ScaledDotProductionAttentionKwargs {
-    training?: boolean;
-    dropout?: number;
-    causal?: boolean;
-    scaling_factor?: number;
-}
-/**
- * This MultiHead Attention layer implements the algorithm as described in
- * the paper "Attention is all you Need" Vaswani et al., 2017.
- *
- * @param numHeads number of attention heads to use
- * @param embedDim the embedding size of the input (input embeddings, typically the last dimension)
- * @param causal use causal masking, default `false`
- * @param dropout use dropout during the attention calculations, default `0.0`
- * @param useBias use bias for the dense sublayers, default `true`
- *
- * The TensorFlow version uses tf.einsum, whose gradient op has not yet been
- * implemented (https://github.com/tensorflow/tfjs/pull/4955#discussion_r619219334),
- * therefore we follow the PyTorch implementation described in:
- * https://docs.pytorch.org/tutorials/intermediate/transformer_building_blocks.html#multiheadattention
- * https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
- *
- * This implementation is different from TensorFlow's whose attention weights
- * are shaped [embed dim, heads, embed dim] where as PyTorch and OpenAI's attention weights
- * are shaped [embed dim, embed dim]
- * https://github.com/pytorch/pytorch/blob/134179474539648ba7dee1317959529fbd0e7f89/torch/nn/modules/activation.py#L1080
- * https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/model.py#L53
- *
- * TODO: implement a fast track for self attention (query = key = value)
- * where a single dense layer combines and replaces the query, key and projection layers
- *
- * TODO: add kDim and vDim to accept key and values whose embedding dimensions differ from query's.
- */
-export declare class MultiHeadAttention extends tf.layers.Layer {
-    static className: string;
-    protected readonly numHeads: number;
-    protected readonly embedDim: number;
-    protected readonly useBias: boolean;
-    protected readonly dropout: number;
-    protected readonly causal: boolean;
-    protected readonly queryProjection: tf.layers.Layer;
-    protected readonly keyProjection: tf.layers.Layer;
-    protected readonly valueProjection: tf.layers.Layer;
-    protected readonly outputProjection: tf.layers.Layer;
-    constructor({ numHeads, embedDim, useBias, dropout, causal, ...args }: MultiHeadAttentionArgs);
-    /**
-     * Forward propagation. Provide one input tensor or three identical tensors to self-attention.
-     * @param inputs a single tensor for self-attention or an array of exactly three
-     *   tensors that are either identical (self-attention) or different (cross-attention)
-     * @param kwargs.packingMask a mask to prevent tokens from attending across document boundaries
-     */
-    call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs & {
-        packingMask?: tf.Tensor;
-        causalMask?: tf.Tensor;
-    }): tf.Tensor | tf.Tensor[];
-    /**
-     * Forward propagation
-     */
-    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 applyInputProjections(query_input: tf.Tensor, key_input: tf.Tensor, value_input: tf.Tensor): {
-        query: tf.Tensor;
-        key: tf.Tensor;
-        value: tf.Tensor;
-    };
-    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;
-    };
-    /**
-     * Applies the scaled dot-product formula: softmax(QK_t / sqrt(d_k))V,
-     * formula (1) of the 2017 paper Attention Is All You Need
-     *
-     * @param attentionMask a mask to prevent tokens from being
-     *   attended to (usually for padding tokens). It should have the shape
-     *   [batch, head, query_sequence_len, key_sequence_len]. To use in
-     *   conjunction with causal masking, the tensor should be a boolean type
-     *   where false indicates a masked token.
-     * @param packingMask a mask to prevent tokens from attending across document boundaries
-     */
-    static scaledDotProductionAttention(query: tf.Tensor, key: tf.Tensor, value: tf.Tensor, attentionMask: tf.Tensor | null, packingMask: tf.Tensor | null, causalMask: tf.Tensor | null, dropout: number, causal: boolean, kwargs?: ScaledDotProductionAttentionKwargs): tf.Tensor;
-    build(inputShape: tf.Shape | tf.Shape[]): void;
-    /**
-     * MultiHead attention's output is the same shape the query's.
-     */
-    computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[];
-    getConfig(): {
-        numHeads: number;
-        embedDim: number;
-        useBias: boolean;
-        causal: boolean;
-        dropout: number;
-        name: string;
-    };
-}
-//# sourceMappingURL=multihead_attention.d.ts.map

package/dist/src/layers/multihead_attention.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"multihead_attention.d.ts","sourceRoot":"","sources":["../../../src/layers/multihead_attention.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;AAIjE,MAAM,WAAW,sBAAuB,SAAQ,SAAS;IACrD,QAAQ,EAAE,MAAM,CAAC;IACjB,QAAQ,EAAE,MAAM,CAAC;IACjB,OAAO,CAAC,EAAE,OAAO,CAAC;IAClB,OAAO,CAAC,EAAE,MAAM,CAAC;IACjB,MAAM,CAAC,EAAE,OAAO,CAAC;CACpB;AAGD,MAAM,WAAW,kCAAkC;IAC/C,QAAQ,CAAC,EAAE,OAAO,CAAC;IACnB,OAAO,CAAC,EAAE,MAAM,CAAC;IACjB,MAAM,CAAC,EAAE,OAAO,CAAC;IACjB,cAAc,CAAC,EAAE,MAAM,CAAC;CAC3B;AAGD;;;;;;;;;;;;;;;;;;;;;;;;;;GA0BG;AACH,qBAAa,kBAAmB,SAAQ,EAAE,CAAC,MAAM,CAAC,KAAK;IACnD,MAAM,CAAC,SAAS,SAAwB;IACxC,SAAS,CAAC,QAAQ,CAAC,QAAQ,EAAE,MAAM,CAAC;IACpC,SAAS,CAAC,QAAQ,CAAC,QAAQ,EAAE,MAAM,CAAC;IACpC,SAAS,CAAC,QAAQ,CAAC,OAAO,EAAE,OAAO,CAAC;IACpC,SAAS,CAAC,QAAQ,CAAC,OAAO,EAAE,MAAM,CAAC;IACnC,SAAS,CAAC,QAAQ,CAAC,MAAM,EAAE,OAAO,CAAC;IAInC,SAAS,CAAC,QAAQ,CAAC,eAAe,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IACpD,SAAS,CAAC,QAAQ,CAAC,aAAa,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IAClD,SAAS,CAAC,QAAQ,CAAC,eAAe,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;IACpD,SAAS,CAAC,QAAQ,CAAC,gBAAgB,EAAE,EAAE,CAAC,MAAM,CAAC,KAAK,CAAC;gBAGzC,EAAE,QAAQ,EAAE,QAAQ,EAAE,OAAc,EAAE,OAAa,EAAE,MAAc,EAAE,GAAG,IAAI,EAAE,EAAE,sBAAsB;IA0BlH;;;;;OAKG;IACM,IAAI,CACT,MAAM,EAAE,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE,EAC/B,MAAM,EAAE,MAAM,GAAG;QACb,WAAW,CAAC,EAAE,EAAE,CAAC,MAAM,CAAC;QACxB,UAAU,CAAC,EAAE,EAAE,CAAC,MAAM,CAAC;KAC1B,GACF,EAAE,CAAC,MAAM,GAAG,EAAE,CAAC,MAAM,EAAE;IA6B1B;;OAEG;IACH,SAAS,CAAC,OAAO,CACb,WAAW,EAAE,EAAE,CAAC,MAAM,EACtB,SAAS,EAAE,EAAE,CAAC,MAAM,EACpB,WAAW,EAAE,EAAE,CAAC,MAAM,EACtB,YAAY,EAAE,EAAE,CAAC,MAAM,GAAG,IAAI,EAC9B,WAAW,EAAE,EAAE,CAAC,MAAM,GAAG,IAAI,EAC7B,MAAM,EAAE,MAAM,GAAG,EAAE,CAAC,MAAM;IA+B9B,SAAS,CAAC,qBAAqB,CAAC,WAAW,EAAE,EAAE,CAAC,MAAM,EAAE,SAAS,EAAE,EAAE,CAAC,MAAM,EAAE,WAAW,EAAE,EAAE,CAAC,MAAM;eAMtC,EAAE,CAAC,MAAM;aACf,EAAE,CAAC,MAAM;eACH,EAAE,CAAC,MAAM;;IAMvE,SAAS,CAAC,UAAU,CAAC,KAAK,EAAE,EAAE,CAAC,MAAM,EAAE,GAAG,EAAE,EAAE,CAAC,MAAM,EAAE,KAAK,EAAE,EAAE,CAAC,MAAM,EAAE,OAAO,EAAE,MAAM,EAAE;qBAShB,EAAE,CAAC,QAAQ;mBACf,EAAE,CAAC,QAAQ;qBACP,EAAE,CAAC,QAAQ;;IAMrF;;;;;;;;;;OAUG;IACH,MAAM,CAAC,4BAA4B,CAC/B,KAAK,EAAE,EAAE,CAAC,MAAM,EAChB,GAAG,EAAE,EAAE,CAAC,MAAM,EACd,KAAK,EAAE,EAAE,CAAC,MAAM,EAChB,aAAa,EAAE,EAAE,CAAC,MAAM,GAAG,IAAI,EAC/B,WAAW,EAAE,EAAE,CAAC,MAAM,GAAG,IAAI,EAC7B,UAAU,EAAE,EAAE,CAAC,MAAM,GAAG,IAAI,EAC5B,OAAO,EAAE,MAAM,EACf,MAAM,EAAE,OAAO,EACf,MAAM,GAAE,kCAAuC,GAChD,EAAE,CAAC,MAAM;IA0EH,KAAK,CAAC,UAAU,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,IAAI;IA4CvD;;OAEG;IACM,kBAAkB,CAAC,UAAU,EAAE,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE,GAAG,EAAE,CAAC,KAAK,GAAG,EAAE,CAAC,KAAK,EAAE;IAK5E,SAAS;;;;;;;;CAgBrB"}