@stellarapp/tfjs-stellar 1.0.0 → 1.0.1

package/{src/models/u_net.ts → dist/models/u_net.js}

@@ -1,145 +1,73 @@
 import * as tf from "@tensorflow/tfjs";
-import { type ActivationIdentifier } from "@tensorflow/tfjs-layers/dist/keras_format/activation_config";
-
-
-export interface UNetArgs {
-    /**
-     * The starting number of filters.
-     */
-    filters: number;
-    /**
-     * The number of categories. For binary segmentation, `units=1`.
-     */
-    units: number;
-    /**
-     * The activation of the final output convolution layer. Defaults to `sigmoid` if `categories=1`, else `softmax`.
-     */
-    activation?: ActivationIdentifier;
-    /**
-     * The depth of the U-Net or the number of contractions and the number of expansions.
-     */
-    depth: number;
-    /**
-     * Adds residual connections to transform the model into a ResUNet. Defaults to `false`.
-     */
-    residual?: boolean;
-    /**
-     * Adds batch normalization to convolutions. Best used for batched inputs. Defaults to `false`.
-     */
-    batchNorm?: boolean;
-    /**
-     * Set the unbatched input shape of the U-Net in the format `[height, width, channels]`. Defaults to `[null, null, 3]`. If set, only channels is mandatory.
-     */
-    inputShape?: [number | null, number | null, number];
-}
-
-
-export type UNetModelArgs = UNetArgs & Omit<tf.SequentialArgs, "layers">;
-
-
 export class UNetModel extends tf.Sequential {
-
-    constructor(args: UNetModelArgs) {
-        const {
-            filters,
-            units,
-            activation = units == 1 ? "sigmoid" : "softmax",
-            depth,
-            residual = false,
-            batchNorm = false,
-            inputShape = [null, null, 3],
-            ...sequentialArgs
-        } = args;
-
+    constructor(args) {
+        const { filters, units, activation = units == 1 ? "sigmoid" : "softmax", depth, residual = false, batchNorm = false, inputShape = [null, null, 3], ...sequentialArgs } = args;
         sequentialArgs.name = sequentialArgs.name ?? "unet_model";
-
         super({
             ...sequentialArgs,
             // calling user should not modify the layers after instantiation
             layers: [createUNet({ filters, units, activation, depth, residual, batchNorm, inputShape })]
         });
     }
-
-
-    override summary(lineLength?: number, positions?: number[], printFn?: (message?: any, ...optionalParams: any[]) => void): void {
+    summary(lineLength, positions, printFn) {
         super.summary(lineLength, positions, printFn);
-        (this.layers[0] as tf.LayersModel).summary(lineLength, positions, printFn);
+        this.layers[0].summary(lineLength, positions, printFn);
     }
 }
-
-
-export function createUNet({ filters, depth, units, activation, residual = false, batchNorm = false, inputShape = [null, null, 3] }: UNetModelArgs) {
+export function createUNet({ filters, depth, units, activation, residual = false, batchNorm = false, inputShape = [null, null, 3] }) {
     if (units < 1) {
         throw Error(`createUNet: units should be >= 1, got ${units}`);
     }
-
     const [image_height, image_width] = inputShape;
     const divisble_by = 2 ** depth;
-
     if ((image_height != null && image_height % divisble_by != 0) ||
         image_width != null && image_width % divisble_by != 0) {
-        throw Error(`createUNet: the input height and width must be divisible by 2^depth (${divisble_by})`)
+        throw Error(`createUNet: the input height and width must be divisible by 2^depth (${divisble_by})`);
     }
-
     const input = tf.input({ shape: inputShape });
-
-    const skip_connection: tf.SymbolicTensor[] = [];
-
+    const skip_connection = [];
     let x = input;
-
     // calculate the filter sizes for each level
     const filter_sizes = Array.from({ length: depth }, (_, i) => filters * (2 ** i));
-
     for (const filter_size of filter_sizes) {
         const contraction = contractionBlock(x, filter_size, residual, batchNorm, `contraction-f${filter_size}`);
-
-        x = tf.layers.maxPooling2d({ poolSize: 2, strides: 2, name: `pool-f${filter_size}` }).apply(contraction) as tf.SymbolicTensor;
+        x = tf.layers.maxPooling2d({ poolSize: 2, strides: 2, name: `pool-f${filter_size}` }).apply(contraction);
         skip_connection.push(contraction);
     }
-
-    x = contractionBlock(x, filter_sizes.at(-1)! * 2, residual, batchNorm, "bottleneck");
-
+    x = contractionBlock(x, filter_sizes.at(-1) * 2, residual, batchNorm, "bottleneck");
     for (let i = filter_sizes.length - 1; i >= 0; i--) {
         x = expansionBlock(x, skip_connection[i], filter_sizes[i], residual, batchNorm, `expansion-f${filter_sizes[i]}`);
     }
-
     const output = tf.layers.conv2d({
         filters: units,
         kernelSize: 1,
         padding: "same",
         activation: activation ?? (units == 1 ? "sigmoid" : "softmax"),
         name: "output-conv"
-    }).apply(x) as tf.SymbolicTensor;
-
+    }).apply(x);
     return tf.model({ inputs: input, outputs: output, name: "u_net" });
 }
-
-
-export async function loadUNetModel(pathOrIOHandler: string | tf.io.IOHandler, options?: tf.io.LoadOptions) {
+export async function loadUNetModel(pathOrIOHandler, options) {
     const model = await tf.loadLayersModel(pathOrIOHandler, options);
     const unet = createUNet({ depth: 1, filters: 4, units: 1 }); // these are dummy args that are overwritten
     const { name, ...rest } = model;
     Object.assign(unet, rest);
-
     return unet;
 }
-
-
 /**
  * The contraction block of a U-Net
- * 
+ *
  * Conv > BN > ReLU > Conv > BN + residual > ReLU
- * 
+ *
  * TODO: for residual, change order to (BN > ReLU > Conv)x2 + residual
- * 
+ *
  * @param x a previous layer's symbolic output
  * @param filters the number of filters, usually half the previous expansion block's
  * @param residual includes a residual connection
  * @param batchNorm applies batch normalization before ReLU activation
  * @param name a unique name for the contraction block
  */
-function contractionBlock(x: tf.SymbolicTensor, filters: number, residual: boolean, batchNorm: boolean, name: string) {
-
+function contractionBlock(x, filters, residual, batchNorm, name) {
     const conv1 = tf.layers.conv2d({
         filters,
         kernelSize: 3,
@@ -149,7 +77,6 @@ function contractionBlock(x: tf.SymbolicTensor, filters: number, residual: boole
         name: `${name}-1-conv2d`
     });
     const relu1 = tf.layers.reLU({ name: `${name}-1-relu` });
-
     const conv2 = tf.layers.conv2d({
         filters,
         kernelSize: 3,
@@ -159,24 +86,17 @@ function contractionBlock(x: tf.SymbolicTensor, filters: number, residual: boole
         name: `${name}-2-conv2d`
     });
     const relu2 = tf.layers.reLU({ name: `${name}-2-relu` });
-
     let forward = conv1.apply(x);
-
     if (batchNorm) {
         forward = tf.layers.batchNormalization({ name: `${name}-1-batchnorm` }).apply(forward);
     }
-
     forward = relu1.apply(forward);
-
     forward = conv2.apply(forward);
-
     if (batchNorm) {
         forward = tf.layers.batchNormalization({ name: `${name}-2-batchnorm` }).apply(forward);
     }
-
     if (residual) {
         let residual_skip = x;
-
         if (x.shape[x.shape.length - 1] != filters) {
             // a 1x1 convolution on the input to ensure the residual connection's
             // channels/filters dim matches the convolution output
@@ -187,32 +107,26 @@ function contractionBlock(x: tf.SymbolicTensor, filters: number, residual: boole
                 useBias: !batchNorm,
                 kernelInitializer: "heNormal",
                 name: `${name}-residual`
-            }).apply(x) as tf.SymbolicTensor;
+            }).apply(x);
         }
-
         if (batchNorm) {
             residual_skip = tf.layers.batchNormalization({
                 name: `${name}-residual-batchnorm`
-            }).apply(residual_skip) as tf.SymbolicTensor;
+            }).apply(residual_skip);
         }
-
         forward = tf.layers.add().apply([
-            residual_skip as tf.SymbolicTensor,
-            forward as tf.SymbolicTensor
-        ])
+            residual_skip,
+            forward
+        ]);
     }
-
     forward = relu2.apply(forward);
-
-    return forward as tf.SymbolicTensor;
+    return forward;
 }
-
-
 /**
  * The expansion block of a U-Net
- * 
+ *
  * Upconv + skip > contraction block
- *  
+ *
  * @param x a previous layer's symbolic output
  * @param skip the corresponding contraction block's output (before pool), shape matches `x`
  * @param filters the number of filters, usually half the previous expansion block's
@@ -220,8 +134,7 @@ function contractionBlock(x: tf.SymbolicTensor, filters: number, residual: boole
  * @param batchNorm apply batch normalization, should be `false` when batch size is `1`
  * @param name a unique name for the contraction block
  */
-function expansionBlock(x: tf.SymbolicTensor, skip: tf.SymbolicTensor, filters: number, residual: boolean, batchNorm: boolean, name: string) {
-
+function expansionBlock(x, skip, filters, residual, batchNorm, name) {
     const upconv = tf.layers.conv2dTranspose({
         filters,
         padding: "same",
@@ -230,11 +143,9 @@ function expansionBlock(x: tf.SymbolicTensor, skip: tf.SymbolicTensor, filters:
         kernelInitializer: "heNormal",
         name: `${name}-upconv`
     });
-
     const concat = tf.layers.concatenate({ axis: -1, name: `${name}-concat-upconv-skip` });
-
-    let forward = upconv.apply(x) as tf.SymbolicTensor;
-    forward = concat.apply([forward, skip]) as tf.SymbolicTensor;
-
+    let forward = upconv.apply(x);
+    forward = concat.apply([forward, skip]);
     return contractionBlock(forward, filters, residual, batchNorm, name);
 }
+//# sourceMappingURL=u_net.js.map

package/dist/models/u_net.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"u_net.js","sourceRoot":"","sources":["../../src/models/u_net.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAE,MAAM,kBAAkB,CAAC;AAuCvC,MAAM,OAAO,SAAU,SAAQ,EAAE,CAAC,UAAU;IAExC,YAAY,IAAmB;QAC3B,MAAM,EACF,OAAO,EACP,KAAK,EACL,UAAU,GAAG,KAAK,IAAI,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,CAAC,SAAS,EAC/C,KAAK,EACL,QAAQ,GAAG,KAAK,EAChB,SAAS,GAAG,KAAK,EACjB,UAAU,GAAG,CAAC,IAAI,EAAE,IAAI,EAAE,CAAC,CAAC,EAC5B,GAAG,cAAc,EACpB,GAAG,IAAI,CAAC;QAET,cAAc,CAAC,IAAI,GAAG,cAAc,CAAC,IAAI,IAAI,YAAY,CAAC;QAE1D,KAAK,CAAC;YACF,GAAG,cAAc;YACjB,gEAAgE;YAChE,MAAM,EAAE,CAAC,UAAU,CAAC,EAAE,OAAO,EAAE,KAAK,EAAE,UAAU,EAAE,KAAK,EAAE,QAAQ,EAAE,SAAS,EAAE,UAAU,EAAE,CAAC,CAAC;SAC/F,CAAC,CAAC;IACP,CAAC;IAGQ,OAAO,CAAC,UAAmB,EAAE,SAAoB,EAAE,OAA2D;QACnH,KAAK,CAAC,OAAO,CAAC,UAAU,EAAE,SAAS,EAAE,OAAO,CAAC,CAAC;QAC7C,IAAI,CAAC,MAAM,CAAC,CAAC,CAAoB,CAAC,OAAO,CAAC,UAAU,EAAE,SAAS,EAAE,OAAO,CAAC,CAAC;IAC/E,CAAC;CACJ;AAGD,MAAM,UAAU,UAAU,CAAC,EAAE,OAAO,EAAE,KAAK,EAAE,KAAK,EAAE,UAAU,EAAE,QAAQ,GAAG,KAAK,EAAE,SAAS,GAAG,KAAK,EAAE,UAAU,GAAG,CAAC,IAAI,EAAE,IAAI,EAAE,CAAC,CAAC,EAAiB;IAC9I,IAAI,KAAK,GAAG,CAAC,EAAE,CAAC;QACZ,MAAM,KAAK,CAAC,yCAAyC,KAAK,EAAE,CAAC,CAAC;IAClE,CAAC;IAED,MAAM,CAAC,YAAY,EAAE,WAAW,CAAC,GAAG,UAAU,CAAC;IAC/C,MAAM,WAAW,GAAG,CAAC,IAAI,KAAK,CAAC;IAE/B,IAAI,CAAC,YAAY,IAAI,IAAI,IAAI,YAAY,GAAG,WAAW,IAAI,CAAC,CAAC;QACzD,WAAW,IAAI,IAAI,IAAI,WAAW,GAAG,WAAW,IAAI,CAAC,EAAE,CAAC;QACxD,MAAM,KAAK,CAAC,wEAAwE,WAAW,GAAG,CAAC,CAAA;IACvG,CAAC;IAED,MAAM,KAAK,GAAG,EAAE,CAAC,KAAK,CAAC,EAAE,KAAK,EAAE,UAAU,EAAE,CAAC,CAAC;IAE9C,MAAM,eAAe,GAAwB,EAAE,CAAC;IAEhD,IAAI,CAAC,GAAG,KAAK,CAAC;IAEd,4CAA4C;IAC5C,MAAM,YAAY,GAAG,KAAK,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,KAAK,EAAE,EAAE,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,OAAO,GAAG,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,CAAC;IAEjF,KAAK,MAAM,WAAW,IAAI,YAAY,EAAE,CAAC;QACrC,MAAM,WAAW,GAAG,gBAAgB,CAAC,CAAC,EAAE,WAAW,EAAE,QAAQ,EAAE,SAAS,EAAE,gBAAgB,WAAW,EAAE,CAAC,CAAC;QAEzG,CAAC,GAAG,EAAE,CAAC,MAAM,CAAC,YAAY,CAAC,EAAE,QAAQ,EAAE,CAAC,EAAE,OAAO,EAAE,CAAC,EAAE,IAAI,EAAE,SAAS,WAAW,EAAE,EAAE,CAAC,CAAC,KAAK,CAAC,WAAW,CAAsB,CAAC;QAC9H,eAAe,CAAC,IAAI,CAAC,WAAW,CAAC,CAAC;IACtC,CAAC;IAED,CAAC,GAAG,gBAAgB,CAAC,CAAC,EAAE,YAAY,CAAC,EAAE,CAAC,CAAC,CAAC,CAAE,GAAG,CAAC,EAAE,QAAQ,EAAE,SAAS,EAAE,YAAY,CAAC,CAAC;IAErF,KAAK,IAAI,CAAC,GAAG,YAAY,CAAC,MAAM,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QAChD,CAAC,GAAG,cAAc,CAAC,CAAC,EAAE,eAAe,CAAC,CAAC,CAAC,EAAE,YAAY,CAAC,CAAC,CAAC,EAAE,QAAQ,EAAE,SAAS,EAAE,cAAc,YAAY,CAAC,CAAC,CAAC,EAAE,CAAC,CAAC;IACrH,CAAC;IAED,MAAM,MAAM,GAAG,EAAE,CAAC,MAAM,CAAC,MAAM,CAAC;QAC5B,OAAO,EAAE,KAAK;QACd,UAAU,EAAE,CAAC;QACb,OAAO,EAAE,MAAM;QACf,UAAU,EAAE,UAAU,IAAI,CAAC,KAAK,IAAI,CAAC,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,CAAC,SAAS,CAAC;QAC9D,IAAI,EAAE,aAAa;KACtB,CAAC,CAAC,KAAK,CAAC,CAAC,CAAsB,CAAC;IAEjC,OAAO,EAAE,CAAC,KAAK,CAAC,EAAE,MAAM,EAAE,KAAK,EAAE,OAAO,EAAE,MAAM,EAAE,IAAI,EAAE,OAAO,EAAE,CAAC,CAAC;AACvE,CAAC;AAGD,MAAM,CAAC,KAAK,UAAU,aAAa,CAAC,eAAyC,EAAE,OAA2B;IACtG,MAAM,KAAK,GAAG,MAAM,EAAE,CAAC,eAAe,CAAC,eAAe,EAAE,OAAO,CAAC,CAAC;IACjE,MAAM,IAAI,GAAG,UAAU,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,OAAO,EAAE,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,CAAC,CAAC,CAAC,4CAA4C;IACzG,MAAM,EAAE,IAAI,EAAE,GAAG,IAAI,EAAE,GAAG,KAAK,CAAC;IAChC,MAAM,CAAC,MAAM,CAAC,IAAI,EAAE,IAAI,CAAC,CAAC;IAE1B,OAAO,IAAI,CAAC;AAChB,CAAC;AAGD;;;;;;;;;;;;GAYG;AACH,SAAS,gBAAgB,CAAC,CAAoB,EAAE,OAAe,EAAE,QAAiB,EAAE,SAAkB,EAAE,IAAY;IAEhH,MAAM,KAAK,GAAG,EAAE,CAAC,MAAM,CAAC,MAAM,CAAC;QAC3B,OAAO;QACP,UAAU,EAAE,CAAC;QACb,OAAO,EAAE,MAAM;QACf,OAAO,EAAE,CAAC,SAAS;QACnB,iBAAiB,EAAE,UAAU;QAC7B,IAAI,EAAE,GAAG,IAAI,WAAW;KAC3B,CAAC,CAAC;IACH,MAAM,KAAK,GAAG,EAAE,CAAC,MAAM,CAAC,IAAI,CAAC,EAAE,IAAI,EAAE,GAAG,IAAI,SAAS,EAAE,CAAC,CAAC;IAEzD,MAAM,KAAK,GAAG,EAAE,CAAC,MAAM,CAAC,MAAM,CAAC;QAC3B,OAAO;QACP,UAAU,EAAE,CAAC;QACb,OAAO,EAAE,MAAM;QACf,OAAO,EAAE,CAAC,SAAS;QACnB,iBAAiB,EAAE,UAAU;QAC7B,IAAI,EAAE,GAAG,IAAI,WAAW;KAC3B,CAAC,CAAC;IACH,MAAM,KAAK,GAAG,EAAE,CAAC,MAAM,CAAC,IAAI,CAAC,EAAE,IAAI,EAAE,GAAG,IAAI,SAAS,EAAE,CAAC,CAAC;IAEzD,IAAI,OAAO,GAAG,KAAK,CAAC,KAAK,CAAC,CAAC,CAAC,CAAC;IAE7B,IAAI,SAAS,EAAE,CAAC;QACZ,OAAO,GAAG,EAAE,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,IAAI,EAAE,GAAG,IAAI,cAAc,EAAE,CAAC,CAAC,KAAK,CAAC,OAAO,CAAC,CAAC;IAC3F,CAAC;IAED,OAAO,GAAG,KAAK,CAAC,KAAK,CAAC,OAAO,CAAC,CAAC;IAE/B,OAAO,GAAG,KAAK,CAAC,KAAK,CAAC,OAAO,CAAC,CAAC;IAE/B,IAAI,SAAS,EAAE,CAAC;QACZ,OAAO,GAAG,EAAE,CAAC,MAAM,CAAC,kBAAkB,CAAC,EAAE,IAAI,EAAE,GAAG,IAAI,cAAc,EAAE,CAAC,CAAC,KAAK,CAAC,OAAO,CAAC,CAAC;IAC3F,CAAC;IAED,IAAI,QAAQ,EAAE,CAAC;QACX,IAAI,aAAa,GAAG,CAAC,CAAC;QAEtB,IAAI,CAAC,CAAC,KAAK,CAAC,CAAC,CAAC,KAAK,CAAC,MAAM,GAAG,CAAC,CAAC,IAAI,OAAO,EAAE,CAAC;YACzC,qEAAqE;YACrE,sDAAsD;YACtD,aAAa,GAAG,EAAE,CAAC,MAAM,CAAC,MAAM,CAAC;gBAC7B,OAAO;gBACP,UAAU,EAAE,CAAC;gBACb,OAAO,EAAE,MAAM;gBACf,OAAO,EAAE,CAAC,SAAS;gBACnB,iBAAiB,EAAE,UAAU;gBAC7B,IAAI,EAAE,GAAG,IAAI,WAAW;aAC3B,CAAC,CAAC,KAAK,CAAC,CAAC,CAAsB,CAAC;QACrC,CAAC;QAED,IAAI,SAAS,EAAE,CAAC;YACZ,aAAa,GAAG,EAAE,CAAC,MAAM,CAAC,kBAAkB,CAAC;gBACzC,IAAI,EAAE,GAAG,IAAI,qBAAqB;aACrC,CAAC,CAAC,KAAK,CAAC,aAAa,CAAsB,CAAC;QACjD,CAAC;QAED,OAAO,GAAG,EAAE,CAAC,MAAM,CAAC,GAAG,EAAE,CAAC,KAAK,CAAC;YAC5B,aAAkC;YAClC,OAA4B;SAC/B,CAAC,CAAA;IACN,CAAC;IAED,OAAO,GAAG,KAAK,CAAC,KAAK,CAAC,OAAO,CAAC,CAAC;IAE/B,OAAO,OAA4B,CAAC;AACxC,CAAC;AAGD;;;;;;;;;;;GAWG;AACH,SAAS,cAAc,CAAC,CAAoB,EAAE,IAAuB,EAAE,OAAe,EAAE,QAAiB,EAAE,SAAkB,EAAE,IAAY;IAEvI,MAAM,MAAM,GAAG,EAAE,CAAC,MAAM,CAAC,eAAe,CAAC;QACrC,OAAO;QACP,OAAO,EAAE,MAAM;QACf,UAAU,EAAE,CAAC;QACb,OAAO,EAAE,CAAC;QACV,iBAAiB,EAAE,UAAU;QAC7B,IAAI,EAAE,GAAG,IAAI,SAAS;KACzB,CAAC,CAAC;IAEH,MAAM,MAAM,GAAG,EAAE,CAAC,MAAM,CAAC,WAAW,CAAC,EAAE,IAAI,EAAE,CAAC,CAAC,EAAE,IAAI,EAAE,GAAG,IAAI,qBAAqB,EAAE,CAAC,CAAC;IAEvF,IAAI,OAAO,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,CAAsB,CAAC;IACnD,OAAO,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,OAAO,EAAE,IAAI,CAAC,CAAsB,CAAC;IAE7D,OAAO,gBAAgB,CAAC,OAAO,EAAE,OAAO,EAAE,QAAQ,EAAE,SAAS,EAAE,IAAI,CAAC,CAAC;AACzE,CAAC"}

package/dist/src/index.d.ts

@@ -0,0 +1,6 @@
+export * as layers from "@/layers";
+export * as models from "@/models";
+export * as losses from "@/losses";
+export * from "@/kv_cache";
+export * from "@/metrics";
+//# sourceMappingURL=index.d.ts.map

package/dist/src/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,cAAc,YAAY,CAAC;AAC3B,cAAc,WAAW,CAAC"}

package/dist/src/index.js

@@ -0,0 +1,6 @@
+export * as layers from "@/layers";
+export * as models from "@/models";
+export * as losses from "@/losses";
+export * from "@/kv_cache";
+export * from "@/metrics";
+//# sourceMappingURL=index.js.map

package/dist/src/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,cAAc,YAAY,CAAC;AAC3B,cAAc,WAAW,CAAC"}

package/dist/src/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/src/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/dist/src/kv_cache.js

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

@@ -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/src/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

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

@@ -0,0 +1 @@
+{"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"}