@genai-fi/nanogpt 0.2.11 → 0.3.0

package/dist/layers/CausalSelfAttention.d.ts

@@ -1,10 +1,10 @@
-import { default as TF } from '@tensorflow/tfjs';
 import { GPTConfig } from '../config';
 import { default as RoPECache } from './RoPECache';
 import { default as BaseLayer } from './BaseLayer';
+import { Tensor, Variable } from '@tensorflow/tfjs-core';
 export type KVCache = {
-    k: TF.Tensor;
-    v: TF.Tensor;
+    k: Tensor;
+    v: Tensor;
     length: number;
     cumulativeLength: number;
 };
@@ -17,25 +17,25 @@ export default class CausalSelfAttention extends BaseLayer {
     private residDropout;
     private bias;
     private maskInf;
-    private tf;
     private divisor;
     private index;
     private _trainable;
     private units;
-    constructor(tf: typeof TF, index: number, config: GPTConfig, ropeCache?: RoPECache | undefined);
+    constructor(index: number, config: GPTConfig, ropeCache?: RoPECache | undefined);
     private build;
-    get variables(): TF.Variable[];
+    get variables(): Variable[];
     get trainable(): boolean;
     set trainable(value: boolean);
-    saveWeights(map: Map<string, TF.Tensor[]>): void;
-    loadWeights(weights: Map<string, TF.Tensor[]>): void;
+    saveWeights(map: Map<string, Tensor[]>): void;
+    loadWeights(weights: Map<string, Tensor[]>): void;
     private getAttentionScores;
     private getAttentionScoresWithPast;
     private getQKV;
     private getOutputProjection;
-    call(x: TF.Tensor, training?: boolean, includeAttention?: boolean, pastKV?: KVCache): {
-        output: TF.Tensor;
-        attention?: TF.Tensor;
+    private updateCache;
+    call(x: Tensor, training?: boolean, includeAttention?: boolean, pastKV?: KVCache): {
+        output: Tensor;
+        attention?: Tensor;
         presentKV?: KVCache;
     };
     dispose(): void;

package/dist/layers/CausalSelfAttention.js

@@ -1,21 +1,31 @@
-import { attentionMask as x } from "../ops/attentionMask.js";
-import j from "./BaseLayer.js";
-import { qkv as w } from "../ops/qkv.js";
-import { rope as y } from "../ops/rope.js";
-class N extends j {
-  constructor(t, i, s, e) {
-    super(), this.ropeCache = e, this.config = s, this.tf = t, this.index = i, this.units = s.nEmbed * 3, this.cProj = this.tf.layers.dense({
-      units: s.nEmbed,
-      useBias: s.biasInLinear,
-      name: `block_${i}_attn_cProj`,
-      kernelInitializer: this.tf.initializers.randomNormal({
+import { attentionMask as C } from "../ops/attentionMask.js";
+import x from "./BaseLayer.js";
+import { qkv as y } from "../ops/qkv.js";
+import { rope as m } from "../ops/rope.js";
+import { appendCache as b } from "../ops/appendCache.js";
+import { w as j, x as f, t as z } from "../index-pWA4_lUh.js";
+import { r as w, l as E, w as D, b as T } from "../random_width-PbCt7RXv.js";
+import { d as L, a as k } from "../exports_layers-7idKoYqh.js";
+import { o as W } from "../ones-Cog-G2ag.js";
+import { z as M } from "../zeros-CCy9C3uU.js";
+import { v as A } from "../variable-Dl_ub3pk.js";
+import { s as g } from "../softmax-Be_lsqUc.js";
+import { m as _ } from "../mat_mul-D7_a4KJn.js";
+import { r as $ } from "../reshape-C8CR_Bad.js";
+class K extends x {
+  constructor(s, t, i) {
+    super(), this.ropeCache = i, this.config = t, this.index = s, this.units = t.nEmbed * 3, this.cProj = L({
+      units: t.nEmbed,
+      useBias: t.biasInLinear,
+      name: `block_${s}_attn_cProj`,
+      kernelInitializer: w({
         mean: 0,
-        stddev: 0.02 / Math.sqrt(2 * s.nLayer)
+        stddev: 0.02 / Math.sqrt(2 * t.nLayer)
       }),
       biasInitializer: "zeros"
-    }), this.attnDropout = this.tf.layers.dropout({ rate: s.dropout }), this.residDropout = this.tf.layers.dropout({ rate: s.dropout }), this.bias = this.tf.linalg.bandPart(this.tf.ones([s.blockSize, s.blockSize]), -1, 0).cast("bool"), this.divisor = 1 / Math.sqrt(s.nEmbed / s.nHead);
-    const o = this.tf.zeros([s.blockSize, s.blockSize]), a = this.tf.fill([s.blockSize, s.blockSize], Number.NEGATIVE_INFINITY);
-    this.maskInf = this.tf.where(this.bias, o, a);
+    }), this.attnDropout = k({ rate: t.dropout }), this.residDropout = k({ rate: t.dropout }), this.bias = E.bandPart(W([t.blockSize, t.blockSize]), -1, 0).cast("bool"), this.divisor = 1 / Math.sqrt(t.nEmbed / t.nHead);
+    const e = M([t.blockSize, t.blockSize]), o = j([t.blockSize, t.blockSize], Number.NEGATIVE_INFINITY);
+    this.maskInf = D(this.bias, e, o);
   }
   config;
   cAttn = null;
@@ -24,14 +34,13 @@ class N extends j {
   residDropout;
   bias;
   maskInf;
-  tf;
   divisor;
   index;
   _trainable = !0;
   units;
   build() {
-    this.cAttn === null && (this.cAttn = this.tf.variable(
-      this.tf.randomNormal([this.config.nEmbed, this.units], 0, 0.02),
+    this.cAttn === null && (this.cAttn = A(
+      T([this.config.nEmbed, this.units], 0, 0.02),
       !0
       //`block_${this.index}_attn_cAttn_kernel`
     ));
@@ -39,70 +48,69 @@ class N extends j {
   get variables() {
     if (this.cAttn === null)
       throw new Error("Layer not built yet");
-    return [this.cAttn, ...this.cProj.trainableWeights.map((t) => t.read())];
+    return [this.cAttn, ...this.cProj.trainableWeights.map((s) => s.read())];
   }
   get trainable() {
     return this._trainable;
   }
-  set trainable(t) {
-    this._trainable = t, this.cAttn && (this.cAttn.trainable = t), this.cProj.trainable = t;
+  set trainable(s) {
+    this._trainable = s, this.cAttn && (this.cAttn.trainable = s), this.cProj.trainable = s;
   }
-  saveWeights(t) {
-    t.set(`block_${this.index}_cAttn`, this.cAttn ? [this.cAttn.clone()] : []), t.set(`block_${this.index}_cProj`, this.cProj.getWeights());
+  saveWeights(s) {
+    s.set(`block_${this.index}_cAttn`, this.cAttn ? [this.cAttn.clone()] : []), s.set(`block_${this.index}_cProj`, this.cProj.getWeights());
   }
-  loadWeights(t) {
-    const i = t.get(`block_${this.index}_cAttn`)?.[0];
-    if (!i) throw new Error(`Weights for block_${this.index}_cAttn not found`);
-    this.cAttn ? this.cAttn.assign(i) : this.cAttn = this.tf.variable(i, !0), this.cProj.setWeights(t.get(`block_${this.index}_cProj`) || []);
+  loadWeights(s) {
+    const t = s.get(`block_${this.index}_cAttn`)?.[0];
+    if (!t) throw new Error(`Weights for block_${this.index}_cAttn not found`);
+    this.cAttn ? this.cAttn.assign(t) : this.cAttn = A(t, !0), this.cProj.setWeights(s.get(`block_${this.index}_cProj`) || []);
   }
-  getAttentionScores(t, i, s) {
-    const e = x(t, i, this.maskInf, this.divisor), o = this.tf.softmax(e, -1);
-    return this.attnDropout.apply(o, { training: s });
+  getAttentionScores(s, t, i) {
+    const e = C(s, t, this.maskInf, this.divisor), o = g(e, -1);
+    return this.attnDropout.apply(o, { training: i });
   }
   // Attention with optional past. If pastLen > 0 and T_cur == 1, no mask needed.
-  getAttentionScoresWithPast(t, i, s, e) {
-    const o = t.shape[2];
-    let r = this.tf.matMul(t, i, !1, !0).mul(this.divisor);
+  getAttentionScoresWithPast(s, t, i, e) {
+    const o = s.shape[2];
+    let n = _(s, t, !1, !0).mul(this.divisor);
     if (o > 1 && e > 0)
       throw new Error("Cannot use past with T_cur > 1");
     if (o > 1) {
-      const c = this.maskInf.slice([0, 0], [o, o]).expandDims(0).expandDims(0);
-      r = r.add(c);
+      const h = this.maskInf.slice([0, 0], [o, o]).expandDims(0).expandDims(0);
+      n = n.add(h);
     }
-    const h = this.tf.softmax(r, -1);
-    return this.attnDropout.apply(h, { training: s });
+    const a = g(n, -1);
+    return this.attnDropout.apply(a, { training: i });
   }
-  getQKV(t) {
-    return w(t, this.cAttn, this.config.nHead);
+  getQKV(s) {
+    return y(s, this.cAttn, this.config.nHead);
   }
-  getOutputProjection(t, i) {
-    const s = t.shape[0], e = t.shape[2], o = this.config.nEmbed, a = t.transpose([0, 2, 1, 3]), r = this.tf.reshape(a, [s, e, o]), h = this.cProj.apply(r);
-    return this.residDropout.apply(h, { training: i });
+  getOutputProjection(s, t) {
+    const i = s.shape[0], e = s.shape[2], o = this.config.nEmbed, r = s.transpose([0, 2, 1, 3]), n = $(r, [i, e, o]), a = this.cProj.apply(n);
+    return this.residDropout.apply(a, { training: t });
+  }
+  updateCache(s, t, i) {
+    const e = this.config.blockSize, o = s.shape[2], r = Math.min(i?.length || 0, e - o), n = i ? b(i.k, s, e) : s, a = i ? b(i.v, t, e) : t;
+    return {
+      k: f(n),
+      v: f(a),
+      length: r + o,
+      cumulativeLength: i ? i.cumulativeLength + o : o
+    };
   }
   // Added optional KV cache support (pastKV). Returns presentKV for chaining.
-  call(t, i = !1, s = !1, e) {
+  call(s, t = !1, i = !1, e) {
     if (e && !this.config.useRope)
       throw new Error("Cannot use pastKV without RoPE enabled");
-    return this.build(), this.tf.tidy(() => {
+    return this.build(), z(() => {
       this.startMemory();
-      const [o, a, r] = this.getQKV(t), h = o.shape[2], c = this.config.blockSize, d = e ? e.cumulativeLength : 0, f = this.ropeCache ? y(o, this.ropeCache, d) : o, m = this.ropeCache ? y(a, this.ropeCache, d) : a;
-      this.ropeCache && (o.dispose(), a.dispose());
-      let n = m, l = r, u = 0;
-      e && (u = e.length, n = this.tf.concat([e.k, m], 2), l = this.tf.concat([e.v, r], 2));
-      const b = n.shape[2];
-      if (b > c) {
-        const k = b - c, A = n.shape[0], g = n.shape[1], _ = n.shape[3];
-        n = n.slice([0, 0, k, 0], [A, g, c, _]), l = l.slice([0, 0, k, 0], [A, g, c, _]), u = c - h;
-      }
-      let p;
-      u > 0 ? p = this.getAttentionScoresWithPast(f, n, i, u) : p = this.getAttentionScores(f, n, i);
-      const P = this.tf.matMul(p, l), S = this.getOutputProjection(P, i), v = {
-        k: this.tf.keep(n),
-        v: this.tf.keep(l),
-        length: u + h,
-        cumulativeLength: e ? e.cumulativeLength + h : h
-      }, I = s ? p.mean(1) : void 0;
-      return this.endMemory("CausalSelfAttention"), { output: S, attention: I, presentKV: v };
+      const [o, r, n] = this.getQKV(s), a = e ? e.cumulativeLength : 0, h = this.ropeCache ? m(o, this.ropeCache, a) : o, p = this.ropeCache ? m(r, this.ropeCache, a) : r;
+      this.ropeCache && (o.dispose(), r.dispose());
+      const u = e ? e.length : 0, l = this.updateCache(p, n, e), d = l.k, v = l.v;
+      e && (p.dispose(), n.dispose());
+      let c;
+      u > 0 ? c = this.getAttentionScoresWithPast(h, d, t, u) : c = this.getAttentionScores(h, d, t);
+      const P = _(c, v), I = this.getOutputProjection(P, t), S = i ? c.mean(1) : void 0;
+      return this.endMemory("CausalSelfAttention"), { output: I, attention: S, presentKV: l };
     });
   }
   dispose() {
@@ -110,5 +118,5 @@ class N extends j {
   }
 }
 export {
-  N as default
+  K as default
 };

package/dist/layers/MLP.d.ts

@@ -1,19 +1,18 @@
-import { default as TF } from '@tensorflow/tfjs';
+import { Tensor, Variable } from '@tensorflow/tfjs-core';
 import { GPTConfig } from '../config';
 import { default as BaseLayer } from './BaseLayer';
 export default class MLP extends BaseLayer {
     private cFc;
     private cProj;
     private dropout;
-    private tf;
     private index;
     private _trainable;
-    constructor(tf: typeof TF, index: number, config: GPTConfig);
-    get variables(): TF.Variable[];
+    constructor(index: number, config: GPTConfig);
+    get variables(): Variable[];
     get trainable(): boolean;
     set trainable(value: boolean);
-    saveWeights(map: Map<string, TF.Tensor[]>): void;
-    loadWeights(weights: Map<string, TF.Tensor[]>): void;
-    call(x: TF.Tensor, training?: boolean): TF.Tensor;
+    saveWeights(map: Map<string, Tensor[]>): void;
+    loadWeights(weights: Map<string, Tensor[]>): void;
+    call(x: Tensor, training?: boolean): Tensor;
     dispose(): void;
 }

package/dist/layers/MLP.js

@@ -1,32 +1,34 @@
-import a from "./BaseLayer.js";
-class l extends a {
+import { t as n } from "../index-pWA4_lUh.js";
+import l from "./BaseLayer.js";
+import { r as s } from "../random_width-PbCt7RXv.js";
+import { d as i, a as c } from "../exports_layers-7idKoYqh.js";
+class u extends l {
   cFc;
   cProj;
   dropout;
-  tf;
   index;
   _trainable = !0;
-  constructor(t, i, e) {
-    super(), this.tf = t, this.index = i, this.cFc = this.tf.layers.dense({
+  constructor(t, e) {
+    super(), this.index = t, this.cFc = i({
       units: e.mlpFactor * e.nEmbed,
       activation: "gelu",
       useBias: e.biasInLinear,
-      kernelInitializer: this.tf.initializers.randomNormal({
+      kernelInitializer: s({
         mean: 0,
         stddev: 0.02
       }),
       biasInitializer: "zeros",
-      name: `block_${i}_mlp_cFc`
-    }), this.cProj = this.tf.layers.dense({
+      name: `block_${t}_mlp_cFc`
+    }), this.cProj = i({
       units: e.nEmbed,
       useBias: e.biasInLinear,
-      kernelInitializer: this.tf.initializers.randomNormal({
+      kernelInitializer: s({
         mean: 0,
         stddev: 0.02 / Math.sqrt(2 * e.nLayer)
       }),
       biasInitializer: "zeros",
-      name: `block_${i}_mlp_cProj`
-    }), this.dropout = this.tf.layers.dropout({ rate: e.dropout });
+      name: `block_${t}_mlp_cProj`
+    }), this.dropout = c({ rate: e.dropout });
   }
   get variables() {
     return [
@@ -46,11 +48,11 @@ class l extends a {
   loadWeights(t) {
     this.cFc.setWeights(t.get(`block_${this.index}_mlpHidden`) || []), this.cProj.setWeights(t.get(`block_${this.index}_mlpOut`) || []);
   }
-  call(t, i = !1) {
-    return this.tf.tidy(() => {
+  call(t, e = !1) {
+    return n(() => {
       this.startMemory();
-      const e = this.cFc.apply(t), s = this.cProj.apply(e), r = this.dropout.apply(s, { training: i });
-      return this.endMemory("MLP"), r;
+      const r = this.cFc.apply(t), a = this.cProj.apply(r), o = this.dropout.apply(a, { training: e });
+      return this.endMemory("MLP"), o;
     });
   }
   dispose() {
@@ -58,5 +60,5 @@ class l extends a {
   }
 }
 export {
-  l as default
+  u as default
 };

package/dist/layers/RMSNorm.d.ts

@@ -1,14 +1,13 @@
-import { default as TF } from '@tensorflow/tfjs';
+import { Tensor, Variable } from '@tensorflow/tfjs-core';
 import { default as BaseLayer } from './BaseLayer';
 export default class RMSNorm extends BaseLayer {
     private gamma;
     private epsilon;
-    private tf;
-    constructor(tf: typeof TF, shape: number[], epsilon?: number, name?: string);
-    get trainableWeights(): TF.Variable[];
+    constructor(shape: number[], epsilon?: number, name?: string);
+    get trainableWeights(): Variable[];
     set trainable(value: boolean);
-    getWeights(): TF.Tensor[];
-    setWeights(weights: TF.Tensor[]): void;
-    apply(x: TF.Tensor): TF.Tensor;
+    getWeights(): Tensor[];
+    setWeights(weights: Tensor[]): void;
+    apply(x: Tensor): Tensor;
     dispose(): void;
 }

package/dist/layers/RMSNorm.js

@@ -1,28 +1,30 @@
+import { t as r } from "../index-pWA4_lUh.js";
 import m from "./BaseLayer.js";
-class o extends m {
+import { v as i } from "../variable-Dl_ub3pk.js";
+import { o } from "../ones-Cog-G2ag.js";
+class d extends m {
   gamma;
   epsilon;
-  tf;
-  constructor(t, s, a = 1e-8, e = "") {
-    super(), this.tf = t, this.epsilon = a, this.gamma = t.variable(t.ones(s), !0, `${e}_gamma`, "float32");
+  constructor(a, s = 1e-8, t = "") {
+    super(), this.epsilon = s, this.gamma = i(o(a), !0, `${t}_gamma`, "float32");
   }
   get trainableWeights() {
     return [this.gamma];
   }
-  set trainable(t) {
-    this.gamma.trainable = t;
+  set trainable(a) {
+    this.gamma.trainable = a;
   }
   getWeights() {
     return [this.gamma];
   }
-  setWeights(t) {
-    this.gamma.assign(t[0]);
+  setWeights(a) {
+    this.gamma.assign(a[0]);
   }
-  apply(t) {
-    return this.tf.tidy(() => {
+  apply(a) {
+    return r(() => {
       this.startMemory();
-      const a = t.square().mean(-1, !0).add(this.epsilon).rsqrt(), r = t.mul(a).mul(this.gamma);
-      return this.endMemory("RMSNorm"), r;
+      const t = a.square().mean(-1, !0).add(this.epsilon).rsqrt(), e = a.mul(t).mul(this.gamma);
+      return this.endMemory("RMSNorm"), e;
     });
   }
   dispose() {
@@ -30,5 +32,5 @@ class o extends m {
   }
 }
 export {
-  o as default
+  d as default
 };

package/dist/layers/RoPECache.d.ts

@@ -1,7 +1,6 @@
-import { default as TF } from '@tensorflow/tfjs';
+import { Tensor } from '@tensorflow/tfjs-core';
 import { GPTConfig } from '../config';
 export default class RoPECache {
-    private readonly tf;
     private readonly config;
     readonly rotaryDim: number;
     private ropeBase;
@@ -9,10 +8,9 @@ export default class RoPECache {
     private ropeCos;
     private ropeSin;
     private ropeCacheLen;
-    constructor(tf: typeof TF, config: GPTConfig);
+    constructor(config: GPTConfig);
     ensureRopeCache(needed: number): void;
-    getCos(): TF.Tensor | null;
-    getSin(): TF.Tensor | null;
-    applyRoPE(q: TF.Tensor, k: TF.Tensor, pastLen: number): [TF.Tensor, TF.Tensor];
+    getCos(): Tensor | null;
+    getSin(): Tensor | null;
     dispose(): void;
 }

package/dist/layers/RoPECache.js

@@ -1,12 +1,36 @@
-class b {
-  constructor(s, r) {
-    this.tf = s, this.config = r;
-    const o = this.config.nEmbed / this.config.nHead;
-    if (this.rotaryDim = o, this.rotaryDim % 2 !== 0)
+import { o as h, h as c, E as f, I as l, f as n, J as m, t as u, x as p } from "../index-pWA4_lUh.js";
+import { c as d, s as C } from "../sin-BJIrfnj7.js";
+import { r as a } from "../range-CcDl05lo.js";
+/**
+ * @license
+ * Copyright 2018 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function x(r) {
+  const s = { x: c(r, "x", "reciprocal") };
+  return f.runKernel(l, s);
+}
+const S = /* @__PURE__ */ h({ reciprocal_: x });
+class y {
+  constructor(o) {
+    this.config = o;
+    const s = this.config.nEmbed / this.config.nHead;
+    if (this.rotaryDim = s, this.rotaryDim % 2 !== 0)
       throw new Error("rotaryDim must be even");
     this.ropeBase = 1e4;
-    const i = this.tf.range(0, this.rotaryDim, 2, "float32"), t = i.div(this.tf.scalar(this.rotaryDim, "float32")), e = this.tf.pow(this.tf.scalar(this.ropeBase, "float32"), t);
-    this.ropeInvFreq = this.tf.reciprocal(e), t.dispose(), e.dispose(), i.dispose(), this.config.useRope === !1 ? (this.ropeCos = null, this.ropeSin = null, this.ropeCacheLen = 0) : this.tf.tidy(() => {
+    const i = a(0, this.rotaryDim, 2, "float32"), e = i.div(n(this.rotaryDim, "float32")), t = m(n(this.ropeBase, "float32"), e);
+    this.ropeInvFreq = S(t), e.dispose(), t.dispose(), i.dispose(), this.config.useRope === !1 ? (this.ropeCos = null, this.ropeSin = null, this.ropeCacheLen = 0) : u(() => {
       this.ensureRopeCache(this.config.blockSize * 4);
     });
   }
@@ -18,11 +42,11 @@ class b {
   ropeSin = null;
   // [cacheLen, rotaryDim/2]
   ropeCacheLen = 0;
-  ensureRopeCache(s) {
-    if (s <= this.ropeCacheLen) return;
+  ensureRopeCache(o) {
+    if (o <= this.ropeCacheLen) return;
     this.ropeCos && this.ropeCos.dispose(), this.ropeSin && this.ropeSin.dispose();
-    const o = this.tf.range(0, s, 1, "float32").expandDims(1).mul(this.ropeInvFreq.expandDims(0));
-    this.ropeCos = this.tf.keep(this.tf.cos(o).expandDims(-1)), this.ropeSin = this.tf.keep(this.tf.sin(o).expandDims(-1)), this.ropeCacheLen = s;
+    const s = Math.max(o, this.ropeCacheLen + this.config.blockSize * 4), e = a(0, s, 1, "float32").expandDims(1).mul(this.ropeInvFreq.expandDims(0));
+    this.ropeCos = p(d(e).expandDims(-1)), this.ropeSin = p(C(e).expandDims(-1)), this.ropeCacheLen = s;
   }
   getCos() {
     return this.ropeCos;
@@ -30,21 +54,10 @@ class b {
   getSin() {
     return this.ropeSin;
   }
-  applyRoPE(s, r, o) {
-    const i = s.shape[3], t = this.rotaryDim;
-    if (t > i) return [s, r];
-    const e = s.shape[2], R = o + e;
-    this.ensureRopeCache(R);
-    const n = t / 2, c = this.ropeCos.slice([o, 0, 0], [e, n, 1]).reshape([1, 1, e, n]), a = this.ropeSin.slice([o, 0, 0], [e, n, 1]).reshape([1, 1, e, n]), h = s.shape[0], p = s.shape[1], f = this.tf.range(0, t, 2, "int32"), l = this.tf.range(1, t, 2, "int32"), d = (u) => {
-      const m = u.slice([0, 0, 0, 0], [h, p, e, t]), C = t < i ? u.slice([0, 0, 0, t], [h, p, e, i - t]) : null, g = this.tf.gather(m, f, 3), D = this.tf.gather(m, l, 3), x = g.mul(c).sub(D.mul(a)), k = D.mul(c).add(g.mul(a)), S = this.tf.stack([x, k], -1).reshape([h, p, e, t]);
-      return C ? this.tf.concat([S, C], 3) : S;
-    }, v = d(s), y = d(r);
-    return f.dispose(), l.dispose(), [v, y];
-  }
   dispose() {
     this.ropeCos && this.ropeCos.dispose(), this.ropeSin && this.ropeSin.dispose(), this.ropeInvFreq.dispose();
   }
 }
 export {
-  b as default
+  y as default
 };

package/dist/layers/TiedEmbedding.d.ts

@@ -1,20 +1,19 @@
-import { default as TF } from '@tensorflow/tfjs';
+import { Tensor, Variable } from '@tensorflow/tfjs-core';
 export default class TiedEmbeddingOutputLayer {
     private vocabSize;
     private embedDim;
-    private tf;
     private tiedWeights;
     private initializer;
-    constructor(tf: typeof TF, config: {
+    constructor(config: {
         vocabSize: number;
         embedDim: number;
         name?: string;
     }, name?: string);
-    get variables(): TF.Variable[];
-    embed(inputs: TF.Tensor): TF.Tensor;
-    project(inputs: TF.Tensor): TF.Tensor;
-    getWeights(): TF.Tensor[];
-    setWeights(weights: TF.Tensor[]): void;
+    get variables(): Variable[];
+    embed(inputs: Tensor): Tensor;
+    project(inputs: Tensor): Tensor;
+    getWeights(): Tensor[];
+    setWeights(weights: Tensor[]): void;
     getConfig(): {
         vocabSize: number;
         embedDim: number;