@genai-fi/nanogpt 0.4.0 → 0.4.2

package/dist/kernel_funcs_utils-CAd1h9X1.js

@@ -0,0 +1,388 @@
+import { aj as E, ak as D, af as B, al as w, n as N, am as v } from "./index-C4JCoBvj.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 C(t) {
+  try {
+    return t.map((e) => E(e));
+  } catch (e) {
+    throw new Error(`Failed to decode encoded string bytes into utf-8, error: ${e}`);
+  }
+}
+function F(t) {
+  return t.map((e) => D(e));
+}
+/**
+ * @license
+ * Copyright 2017 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 R(t) {
+  if (t <= 1)
+    return "int";
+  if (t === 2)
+    return "ivec2";
+  if (t === 3)
+    return "ivec3";
+  if (t === 4)
+    return "ivec4";
+  if (t === 5)
+    return "ivec5";
+  if (t === 6)
+    return "ivec6";
+  throw Error(`GPU for rank ${t} is not yet supported`);
+}
+/**
+ * @license
+ * Copyright 2017 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 y(t) {
+  return B().getBool("WEBGL_USE_SHAPES_UNIFORMS") && t <= 4;
+}
+/**
+ * @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 k(t, e) {
+  return ["x", "y", "z", "w", "u", "v"].slice(0, e).map((a) => `${t}.${a}`);
+}
+function _(t, e) {
+  return e === 1 ? [t] : k(t, e);
+}
+/**
+ * @license
+ * Copyright 2017 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.
+ * =============================================================================
+ */
+class A {
+  constructor(e, a, u) {
+    this.variableNames = ["A", "B"], this.outputShape = w(a, u), this.enableShapeUniforms = y(this.outputShape.length), this.userCode = `
+      float binaryOperation(float a, float b) {
+        ${e}
+      }
+
+      void main() {
+        float a = getAAtOutCoords();
+        float b = getBAtOutCoords();
+        setOutput(binaryOperation(a, b));
+      }
+    `;
+  }
+}
+/**
+ * @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.
+ * =============================================================================
+ */
+class z {
+  constructor(e, a, u, d = !1) {
+    this.variableNames = ["A", "B"], this.supportsBroadcasting = !0, this.packedInputs = !0, this.packedOutput = !0, this.outputShape = w(a, u);
+    const o = this.outputShape.length;
+    this.enableShapeUniforms = y(o);
+    let n = "";
+    if (d)
+      if (o === 0 || N(this.outputShape) === 1)
+        n = `
+          result.y = 0.;
+          result.z = 0.;
+          result.w = 0.;
+        `;
+      else if (n = `
+          ${R(o)} coords = getOutputCoords();
+        `, o === 1)
+        this.enableShapeUniforms ? n += `
+            result.y = (coords + 1) >= outShape ? 0. : result.y;
+            result.z = 0.;
+            result.w = 0.;
+          ` : n += `
+            result.y = (coords + 1) >= ${this.outputShape[0]} ? 0. : result.y;
+            result.z = 0.;
+            result.w = 0.;
+          `;
+      else {
+        const s = _("coords", o);
+        this.enableShapeUniforms ? n += `
+            bool nextRowOutOfBounds =
+              (${s[o - 2]} + 1) >= outShape[${o} - 2];
+            bool nextColOutOfBounds =
+              (${s[o - 1]} + 1) >= outShape[${o} - 1];
+            result.y = nextColOutOfBounds ? 0. : result.y;
+            result.z = nextRowOutOfBounds ? 0. : result.z;
+            result.w = nextColOutOfBounds || nextRowOutOfBounds ? 0. : result.w;
+          ` : n += `
+            bool nextRowOutOfBounds =
+              (${s[o - 2]} + 1) >= ${this.outputShape[o - 2]};
+            bool nextColOutOfBounds =
+              (${s[o - 1]} + 1) >= ${this.outputShape[o - 1]};
+            result.y = nextColOutOfBounds ? 0. : result.y;
+            result.z = nextRowOutOfBounds ? 0. : result.z;
+            result.w = nextColOutOfBounds || nextRowOutOfBounds ? 0. : result.w;
+          `;
+      }
+    this.userCode = `
+      vec4 binaryOperation(vec4 a, vec4 b) {
+        ${e}
+      }
+
+      void main() {
+        vec4 a = getAAtOutCoords();
+        vec4 b = getBAtOutCoords();
+
+        vec4 result = binaryOperation(a, b);
+        ${n}
+
+        setOutput(result);
+      }
+    `;
+  }
+}
+/**
+ * @license
+ * Copyright 2020 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 P(t) {
+  const { inputs: e, backend: a } = t, { x: u } = e;
+  return a.incRef(u.dataId), { dataId: u.dataId, shape: u.shape, dtype: u.dtype };
+}
+/**
+ * @license
+ * Copyright 2020 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 G(t) {
+  const { inputs: e, backend: a } = t, { real: u, imag: d } = e, o = a.makeTensorInfo(u.shape, "complex64"), n = a.texData.get(o.dataId), l = P({ inputs: { x: u }, backend: a }), s = P({ inputs: { x: d }, backend: a });
+  return n.complexTensorInfos = { real: l, imag: s }, o;
+}
+/**
+ * @license
+ * Copyright 2017 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.
+ * =============================================================================
+ */
+class V {
+  constructor(e, a) {
+    this.variableNames = ["A"], this.outputShape = e, this.enableShapeUniforms = y(this.outputShape.length), this.userCode = `
+      float unaryOperation(float x) {
+        ${a}
+      }
+
+      void main() {
+        float x = getAAtOutCoords();
+        float y = unaryOperation(x);
+
+        setOutput(y);
+      }
+    `;
+  }
+}
+const H = "if (isnan(x)) return x;";
+/**
+ * @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.
+ * =============================================================================
+ */
+class L {
+  constructor(e, a) {
+    this.variableNames = ["A"], this.packedInputs = !0, this.packedOutput = !0, this.outputShape = e, this.enableShapeUniforms = y(this.outputShape.length), this.userCode = `
+      vec4 unaryOperation(vec4 x) {
+        ${a}
+      }
+
+      void main() {
+        vec4 x = getAAtOutCoords();
+        vec4 y = unaryOperation(x);
+
+        setOutput(y);
+      }
+    `;
+  }
+}
+/**
+ * @license
+ * Copyright 2020 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 K({ opSnippet: t, packedOpSnippet: e, cpuKernelImpl: a, dtype: u }) {
+  return ({ inputs: d, backend: o }) => {
+    const { x: n } = d, l = o, s = u || n.dtype;
+    if (l.shouldExecuteOnCPU([n]) && a != null) {
+      const p = l.texData.get(n.dataId), x = a(p.values, s);
+      return l.makeTensorInfo(n.shape, s, x);
+    }
+    const i = B().getBool("WEBGL_PACK_UNARY_OPERATIONS") && e != null;
+    let r;
+    return i ? r = new L(n.shape, e) : r = new V(n.shape, t), l.runWebGLProgram(r, [n], s);
+  };
+}
+function Y({ opSnippet: t, packedOpSnippet: e, checkOutOfBounds: a = !1, supportsComplex: u = !1, cpuKernelImpl: d, dtype: o }) {
+  return ({ inputs: n, backend: l }) => {
+    const { a: s, b: i } = n, r = l;
+    if (u && s.dtype === "complex64") {
+      const h = r.texData.get(s.dataId), f = r.texData.get(i.dataId), [g, m] = [
+        [h.complexTensorInfos.real, f.complexTensorInfos.real],
+        [h.complexTensorInfos.imag, f.complexTensorInfos.imag]
+      ].map((S) => {
+        const [c, O] = S, $ = {
+          dataId: c.dataId,
+          dtype: c.dtype,
+          shape: s.shape
+        }, T = {
+          dataId: O.dataId,
+          dtype: O.dtype,
+          shape: i.shape
+        }, U = new A(t, s.shape, i.shape);
+        return r.runWebGLProgram(U, [$, T], v(c.dtype, O.dtype));
+      }), I = G({ inputs: { real: g, imag: m }, backend: r });
+      return r.disposeIntermediateTensorInfo(g), r.disposeIntermediateTensorInfo(m), I;
+    }
+    const p = o || v(s.dtype, i.dtype);
+    if ((s.dtype === "string" || i.dtype === "string" || r.shouldExecuteOnCPU([s, i])) && d != null) {
+      const h = r.texData.get(s.dataId).values, f = r.texData.get(i.dataId).values, g = s.dtype === "string" ? (
+        // tslint:disable-next-line: no-any
+        C(h)
+      ) : h, m = s.dtype === "string" ? (
+        // tslint:disable-next-line: no-any
+        C(f)
+      ) : f, [I, S] = d(s.shape, i.shape, g, m, p), c = r.makeTensorInfo(S, p), O = r.texData.get(c.dataId);
+      return O.values = I, c;
+    }
+    const x = B().getBool("WEBGL_PACK_BINARY_OPERATIONS") && e != null;
+    let b;
+    return x ? b = new z(e, s.shape, i.shape, a) : b = new A(t, s.shape, i.shape), r.runWebGLProgram(b, [s, i], p);
+  };
+}
+export {
+  H as C,
+  F as a,
+  y as b,
+  k as c,
+  Y as d,
+  C as f,
+  R as g,
+  K as u
+};

package/dist/layers/CausalSelfAttention.js

@@ -1,17 +1,18 @@
-import { attentionMask as I } from "../ops/attentionMask.js";
-import y from "./BaseLayer.js";
-import { qkv as z } from "../ops/qkv.js";
-import { rope as P } from "../ops/rope.js";
+import { attentionMask as P } from "../ops/attentionMask.js";
+import T from "./BaseLayer.js";
+import { qkv as y } from "../ops/qkv.js";
+import { rope as w } from "../ops/rope.js";
 import { appendCache as E } from "../ops/appendCache.js";
-import { D as $, F as _, t as x, c as L, e as v, H as W } from "../index-XjBAhiFO.js";
-import { fusedSoftmax as S } from "../ops/fusedSoftmax.js";
-import { l as M, w as O, r as T, d as N, a as U } from "../tfjs_backend-Bzl2SrRo.js";
-import { o as q } from "../ones-BO01zpJG.js";
-import { z as B } from "../zeros-CmJFiC84.js";
-import { v as g } from "../variable-BS4AKqNU.js";
-import { m as C } from "../mat_mul-CKwFEV1Q.js";
-import { r as D } from "../reshape-DFzh97Sc.js";
-class nt extends y {
+import { D as z, F as S, t as $, c as L, e as j, H as O } from "../index-C4JCoBvj.js";
+import { fusedSoftmax as _ } from "../ops/fusedSoftmax.js";
+import { l as W, w as M, d as x } from "../tfjs_backend-Cug-PH75.js";
+import { o as N } from "../ones-Bf3YR48P.js";
+import { z as q } from "../zeros-dnQxFgAD.js";
+import { v as k } from "../variable-LJT9Ld63.js";
+import { r as C, d as I } from "../dropout-DfDdklfL.js";
+import { r as B } from "../reshape-Boe4DuIO.js";
+import { m as F } from "../mat_mul-415y5Qn2.js";
+class nt extends T {
   cAttn = null;
   cProj = null;
   bias;
@@ -22,17 +23,17 @@ class nt extends y {
   units;
   projUnits;
   constructor(t, s) {
-    super(s), this.index = t, this.units = s.gpt.nEmbed * 3, this.projUnits = s.gpt.nEmbed, this.bias = M.bandPart(q([s.gpt.blockSize, s.gpt.blockSize]), -1, 0).cast("bool"), this.divisor = 1 / Math.sqrt(s.gpt.nEmbed / s.gpt.nHead);
-    const e = B([s.gpt.blockSize, s.gpt.blockSize]), n = $([s.gpt.blockSize, s.gpt.blockSize], Number.NEGATIVE_INFINITY);
-    this.maskInf = O(this.bias, e, n);
+    super(s), this.index = t, this.units = s.gpt.nEmbed * 3, this.projUnits = s.gpt.nEmbed, this.bias = W.bandPart(N([s.gpt.blockSize, s.gpt.blockSize]), -1, 0).cast("bool"), this.divisor = 1 / Math.sqrt(s.gpt.nEmbed / s.gpt.nHead);
+    const e = q([s.gpt.blockSize, s.gpt.blockSize]), i = z([s.gpt.blockSize, s.gpt.blockSize], Number.NEGATIVE_INFINITY);
+    this.maskInf = M(this.bias, e, i);
   }
   build() {
-    this.cAttn === null && (this.cAttn = g(
-      T([this.config.gpt.nEmbed, this.units], 0, 0.02),
+    this.cAttn === null && (this.cAttn = k(
+      C([this.config.gpt.nEmbed, this.units], 0, 0.02),
       !0
       //`block_${this.index}_attn_cAttn_kernel`
-    )), this.cProj === null && (this.cProj = g(
-      T([this.projUnits, this.config.gpt.nEmbed], 0, 0.02),
+    )), this.cProj === null && (this.cProj = k(
+      C([this.projUnits, this.config.gpt.nEmbed], 0, 0.02),
       !0
       //`block_${this.index}_attn_cProj_kernel`
     ));
@@ -55,57 +56,54 @@ class nt extends y {
     const s = t.get(`block_${this.index}_cAttn`)?.[0], e = t.get(`block_${this.index}_cProj`)?.[0];
     if (!s) throw new Error(`Weights for block_${this.index}_cAttn not found`);
     if (!e) throw new Error(`Weights for block_${this.index}_cProj not found`);
-    this.cAttn ? this.cAttn.assign(s) : this.cAttn = g(s, !0), this.cProj ? this.cProj.assign(e) : this.cProj = g(e, !0);
+    this.cAttn ? this.cAttn.assign(s) : this.cAttn = k(s, !0), this.cProj ? this.cProj.assign(e) : this.cProj = k(e, !0);
   }
-  getAttentionScores(t, s, e, n) {
-    const o = I(t, s, this.maskInf, this.divisor);
-    return S(o, e ? this.config.gpt.dropout : 0, n);
+  getAttentionScores(t, s, e, i) {
+    const o = P(t, s, this.divisor, this.maskInf);
+    return _(o, e ? this.config.gpt.dropout : 0, i);
   }
   // Attention with optional past. If pastLen > 0 and T_cur == 1, no mask needed.
-  getAttentionScoresWithPast(t, s, e, n, o) {
-    const i = t.shape[2];
-    let r = C(t, s, !1, !0).mul(this.divisor);
-    if (i > 1 && n > 0)
-      throw new Error("Cannot use past with T_cur > 1");
-    if (i > 1) {
-      const c = this.maskInf.slice([0, 0], [i, i]).expandDims(0).expandDims(0);
-      r = r.add(c);
-    }
-    return S(r, e ? this.config.gpt.dropout : 0, o);
+  getAttentionScoresWithPast(t, s, e) {
+    const i = P(t, s, this.divisor, void 0, e);
+    return _(i, 0, 0);
   }
   getQKV(t) {
-    return z(t, this.cAttn, this.config.gpt.nHead);
+    return y(t, this.cAttn, this.config.gpt.nHead);
   }
   getOutputProjection(t) {
-    const s = t.shape[0], e = t.shape[2], n = this.config.gpt.nEmbed, o = t.transpose([0, 2, 1, 3]), i = D(o, [s, e, n]);
-    return N(i, this.cProj);
+    const s = t.shape[0], e = t.shape[2], i = this.config.gpt.nEmbed, o = t.transpose([0, 2, 1, 3]), n = B(o, [s, e, i]);
+    return x(n, this.cProj);
   }
   updateCache(t, s, e) {
-    const n = this.config.gpt.blockSize, o = t.shape[2], i = Math.min(e?.length || 0, n - o), a = e ? E(e.k, t, n) : t, r = e ? E(e.v, s, n) : s;
+    const i = this.config.gpt.blockSize, o = t.shape[2], n = Math.min(e?.length || 0, i - o), r = e ? E(e.k, t, i) : t, a = e ? E(e.v, s, i) : s;
     return {
-      k: _(a),
-      v: _(r),
-      length: i + o,
+      k: S(r),
+      v: S(a),
+      length: n + o,
       cumulativeLength: e ? e.cumulativeLength + o : o
     };
   }
-  forward(t, s = !1, e, n = !1, o) {
-    return x(() => {
+  forward(t, s = !1, e, i = !1, o) {
+    return $(() => {
       this.startMemory();
-      const [i, a, r] = this.getQKV(t), c = o ? o.cumulativeLength : 0, h = this.config.layerConfig.ropeCache, d = h ? P(i, h, c) : i, p = h ? P(a, h, c) : a;
-      h && (i.dispose(), a.dispose());
-      const f = o ? o.length : 0, l = this.updateCache(p, r, o), m = l.k, b = l.v;
-      o && (p.dispose(), r.dispose());
-      let u;
-      f > 0 ? u = this.getAttentionScoresWithPast(d, m, s, f, e) : u = this.getAttentionScores(d, m, s, e);
-      const k = C(u, b), A = this.getOutputProjection(k), w = n ? u.mean(1) : void 0;
-      return this.endMemory("CausalSelfAttention"), { output: A, attention: w, presentKV: l };
+      const [n, r, a] = this.getQKV(t), p = o ? o.cumulativeLength : 0, c = this.config.layerConfig.ropeCache, u = c ? w(n, c, p) : n, f = c ? w(r, c, p) : r;
+      c && (n.dispose(), r.dispose());
+      const g = o ? o.length : 0, d = this.updateCache(f, a, o), l = d.k, m = d.v;
+      o && (f.dispose(), a.dispose());
+      let h;
+      g > 0 ? h = this.getAttentionScoresWithPast(u, l, g) : h = this.getAttentionScores(u, l, s, e), u.dispose(), s && l.dispose();
+      const b = F(h, m);
+      i || h.dispose(), s && m.dispose();
+      const A = this.getOutputProjection(b);
+      b.dispose();
+      const v = i ? h.mean(1) : void 0;
+      return this.endMemory("CausalSelfAttention"), { output: A, attention: v, presentKV: s ? void 0 : d };
     });
   }
-  call(t, s = !1, e = !1, n) {
-    if (n && !this.config.gpt.useRope)
+  call(t, s = !1, e = !1, i) {
+    if (i && !this.config.gpt.useRope)
       throw new Error("Cannot use pastKV without RoPE enabled");
-    if (s && n)
+    if (s && i)
       throw new Error("Cannot use pastKV during training");
     if (t.shape.length !== 3)
       throw new Error(`Input tensor must be rank 3 [B, T, C], got shape ${t.shape}`);
@@ -114,30 +112,33 @@ class nt extends y {
     this.build();
     const o = Math.random() * 1e9;
     if (s && this.config.layerConfig.checkpointAttention) {
-      const a = L(
+      const r = L(
         // @ts-expect-error Invalid params
-        (r, c, h, d) => {
-          const p = this.forward(r, !0, o);
-          p.presentKV?.k.dispose(), p.presentKV?.v.dispose(), d([r]);
-          const f = (l, m) => {
-            const [b] = m, u = v().state.activeTape;
-            v().state.activeTape = [];
-            const k = W((A, w, H) => {
-              const j = this.forward(A, !0, o);
-              return j.presentKV?.k.dispose(), j.presentKV?.v.dispose(), j.output;
-            })([b, c, h], l);
-            return v().state.activeTape = u, k;
+        (a, p, c, u) => {
+          const f = this.forward(a, !0, o);
+          u([a]);
+          const g = (d, l) => {
+            const [m] = l, h = j().state.activeTape;
+            j().state.activeTape = [];
+            const b = O((A, v, R) => this.forward(A, !0, o).output)([m, p, c], d);
+            return j().state.activeTape = h, b;
           };
-          return { value: p.output, gradFunc: f };
+          return { value: f.output, gradFunc: g };
         }
       )(t, this.cAttn, this.cProj);
       if (this.config.gpt.dropout > 0) {
-        const r = U(a, this.config.gpt.dropout);
-        return a.dispose(), { output: r };
+        const a = I(r, this.config.gpt.dropout);
+        return r.dispose(), { output: a };
+      } else
+        return { output: r };
+    } else {
+      const n = this.forward(t, s, o, e, i);
+      if (this.config.gpt.dropout > 0) {
+        const r = I(n.output, this.config.gpt.dropout);
+        return n.output.dispose(), { output: r, attention: n.attention, presentKV: n.presentKV };
       } else
-        return { output: a };
-    } else
-      return this.forward(t, s, o, e, n);
+        return n;
+    }
   }
   dispose() {
     this.cAttn?.dispose(), this.cProj?.dispose(), this.bias.dispose(), this.maskInf.dispose();

package/dist/layers/MLP.d.ts

@@ -3,15 +3,17 @@ import { default as BaseLayer, GPTLayerConfig } from './BaseLayer';
 export default class MLP extends BaseLayer {
     private cFc;
     private cProj;
-    private dropout;
     private index;
     private _trainable;
+    private hiddenUnits;
     constructor(index: number, config: GPTLayerConfig);
+    private build;
     get variables(): Variable[];
     get trainable(): boolean;
     set trainable(value: boolean);
     saveWeights(map: Map<string, Tensor[]>): void;
     loadWeights(weights: Map<string, Tensor[]>): void;
+    forward(x: Tensor): Tensor;
     call(x: Tensor, training?: boolean): Tensor;
     dispose(): void;
 }