@genai-fi/nanogpt 0.3.1 → 0.4.0

package/dist/layers/BaseLayer.d.ts

@@ -1,8 +1,20 @@
+import { GPTConfig } from '../config';
 import { default as MemoryProfiler } from '../utilities/profile';
+import { default as RoPECache } from './RoPECache';
+export interface LayerConfig {
+    checkpointAttention?: boolean;
+    checkpointMLP?: boolean;
+    profiler?: MemoryProfiler;
+    ropeCache?: RoPECache;
+}
+export interface GPTLayerConfig {
+    gpt: GPTConfig;
+    layerConfig: LayerConfig;
+}
 export default abstract class BaseLayer {
-    protected _profiler?: MemoryProfiler;
+    readonly config: GPTLayerConfig;
+    constructor(config: GPTLayerConfig);
     getProfiler(): MemoryProfiler | undefined;
-    setProfiler(value: MemoryProfiler | undefined): void;
     startMemory(): void;
     endMemory(label: string): void;
 }

package/dist/layers/BaseLayer.js

@@ -1,18 +1,18 @@
-class t {
-  _profiler;
-  getProfiler() {
-    return this._profiler;
+class o {
+  config;
+  constructor(r) {
+    this.config = r;
   }
-  setProfiler(r) {
-    this._profiler = r;
+  getProfiler() {
+    return this.config.layerConfig.profiler;
   }
   startMemory() {
-    this._profiler?.startMemory();
+    this.config.layerConfig.profiler?.startMemory();
   }
   endMemory(r) {
-    this._profiler?.endMemory(r);
+    this.config.layerConfig.profiler?.endMemory(r);
   }
 }
 export {
-  t as default
+  o as default
 };

package/dist/layers/CausalSelfAttention.d.ts

@@ -1,6 +1,4 @@
-import { GPTConfig } from '../config';
-import { default as RoPECache } from './RoPECache';
-import { default as BaseLayer } from './BaseLayer';
+import { default as BaseLayer, GPTLayerConfig } from './BaseLayer';
 import { Tensor, Variable } from '@tensorflow/tfjs-core';
 export type KVCache = {
     k: Tensor;
@@ -9,19 +7,16 @@ export type KVCache = {
     cumulativeLength: number;
 };
 export default class CausalSelfAttention extends BaseLayer {
-    private readonly ropeCache?;
-    private config;
     private cAttn;
     private cProj;
-    private attnDropout;
-    private residDropout;
     private bias;
     private maskInf;
     private divisor;
     private index;
     private _trainable;
     private units;
-    constructor(index: number, config: GPTConfig, ropeCache?: RoPECache | undefined);
+    private projUnits;
+    constructor(index: number, config: GPTLayerConfig);
     private build;
     get variables(): Variable[];
     get trainable(): boolean;
@@ -33,6 +28,7 @@ export default class CausalSelfAttention extends BaseLayer {
     private getQKV;
     private getOutputProjection;
     private updateCache;
+    private forward;
     call(x: Tensor, training?: boolean, includeAttention?: boolean, pastKV?: KVCache): {
         output: Tensor;
         attention?: Tensor;

package/dist/layers/CausalSelfAttention.js

@@ -1,122 +1,148 @@
-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-oeUIlUZj.js";
-import { d as L, a as k } from "../exports_layers-tbTBcwMM.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 * t.nLayer)
-      }),
-      biasInitializer: "zeros"
-    }), 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;
+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 { 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 {
   cAttn = null;
-  cProj;
-  attnDropout;
-  residDropout;
+  cProj = null;
   bias;
   maskInf;
   divisor;
   index;
   _trainable = !0;
   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);
+  }
   build() {
-    this.cAttn === null && (this.cAttn = A(
-      T([this.config.nEmbed, this.units], 0, 0.02),
+    this.cAttn === null && (this.cAttn = g(
+      T([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),
+      !0
+      //`block_${this.index}_attn_cProj_kernel`
     ));
   }
   get variables() {
     if (this.cAttn === null)
       throw new Error("Layer not built yet");
-    return [this.cAttn, ...this.cProj.trainableWeights.map((s) => s.read())];
+    return [this.cAttn, this.cProj];
   }
   get trainable() {
     return this._trainable;
   }
-  set trainable(s) {
-    this._trainable = s, this.cAttn && (this.cAttn.trainable = s), this.cProj.trainable = s;
+  set trainable(t) {
+    this._trainable = t, this.cAttn && (this.cAttn.trainable = t), this.cProj && (this.cProj.trainable = t);
   }
-  saveWeights(s) {
-    s.set(`block_${this.index}_cAttn`, this.cAttn ? [this.cAttn.clone()] : []), s.set(`block_${this.index}_cProj`, this.cProj.getWeights());
+  saveWeights(t) {
+    t.set(`block_${this.index}_cAttn`, this.cAttn ? [this.cAttn.clone()] : []), t.set(`block_${this.index}_cProj`, this.cProj ? [this.cProj.clone()] : []);
   }
-  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`) || []);
+  loadWeights(t) {
+    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);
   }
-  getAttentionScores(s, t, i) {
-    const e = C(s, t, this.maskInf, this.divisor), o = g(e, -1);
-    return this.attnDropout.apply(o, { training: i });
+  getAttentionScores(t, s, e, n) {
+    const o = I(t, s, this.maskInf, this.divisor);
+    return S(o, e ? this.config.gpt.dropout : 0, n);
   }
   // Attention with optional past. If pastLen > 0 and T_cur == 1, no mask needed.
-  getAttentionScoresWithPast(s, t, i, e) {
-    const o = s.shape[2];
-    let n = _(s, t, !1, !0).mul(this.divisor);
-    if (o > 1 && e > 0)
+  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 (o > 1) {
-      const h = this.maskInf.slice([0, 0], [o, o]).expandDims(0).expandDims(0);
-      n = n.add(h);
+    if (i > 1) {
+      const c = this.maskInf.slice([0, 0], [i, i]).expandDims(0).expandDims(0);
+      r = r.add(c);
     }
-    const a = g(n, -1);
-    return this.attnDropout.apply(a, { training: i });
+    return S(r, e ? this.config.gpt.dropout : 0, o);
   }
-  getQKV(s) {
-    return y(s, this.cAttn, this.config.nHead);
+  getQKV(t) {
+    return z(t, this.cAttn, this.config.gpt.nHead);
   }
-  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 });
+  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);
   }
-  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;
+  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;
     return {
-      k: f(n),
-      v: f(a),
-      length: r + o,
-      cumulativeLength: i ? i.cumulativeLength + o : o
+      k: _(a),
+      v: _(r),
+      length: i + o,
+      cumulativeLength: e ? e.cumulativeLength + o : o
     };
   }
-  // Added optional KV cache support (pastKV). Returns presentKV for chaining.
-  call(s, t = !1, i = !1, e) {
-    if (e && !this.config.useRope)
-      throw new Error("Cannot use pastKV without RoPE enabled");
-    return this.build(), z(() => {
+  forward(t, s = !1, e, n = !1, o) {
+    return x(() => {
       this.startMemory();
-      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 };
+      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 };
     });
   }
+  call(t, s = !1, e = !1, n) {
+    if (n && !this.config.gpt.useRope)
+      throw new Error("Cannot use pastKV without RoPE enabled");
+    if (s && n)
+      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}`);
+    if (t.shape[2] !== this.config.gpt.nEmbed)
+      throw new Error(`Input tensor last dimension must be ${this.config.gpt.nEmbed}, got ${t.shape[2]}`);
+    this.build();
+    const o = Math.random() * 1e9;
+    if (s && this.config.layerConfig.checkpointAttention) {
+      const a = 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;
+          };
+          return { value: p.output, gradFunc: f };
+        }
+      )(t, this.cAttn, this.cProj);
+      if (this.config.gpt.dropout > 0) {
+        const r = U(a, this.config.gpt.dropout);
+        return a.dispose(), { output: r };
+      } else
+        return { output: a };
+    } else
+      return this.forward(t, s, o, e, n);
+  }
   dispose() {
-    this.cAttn?.dispose(), this.cProj.dispose(), this.attnDropout.dispose(), this.residDropout.dispose(), this.bias.dispose(), this.maskInf.dispose();
+    this.cAttn?.dispose(), this.cProj?.dispose(), this.bias.dispose(), this.maskInf.dispose();
   }
 }
 export {
-  K as default
+  nt as default
 };

package/dist/layers/MLP.d.ts

@@ -1,13 +1,12 @@
 import { Tensor, Variable } from '@tensorflow/tfjs-core';
-import { GPTConfig } from '../config';
-import { default as BaseLayer } from './BaseLayer';
+import { default as BaseLayer, GPTLayerConfig } from './BaseLayer';
 export default class MLP extends BaseLayer {
     private cFc;
     private cProj;
     private dropout;
     private index;
     private _trainable;
-    constructor(index: number, config: GPTConfig);
+    constructor(index: number, config: GPTLayerConfig);
     get variables(): Variable[];
     get trainable(): boolean;
     set trainable(value: boolean);

package/dist/layers/MLP.js

@@ -1,64 +1,7 @@
-import { t as n } from "../index-pWA4_lUh.js";
-import l from "./BaseLayer.js";
-import { r as s } from "../random_width-oeUIlUZj.js";
-import { d as i, a as c } from "../exports_layers-tbTBcwMM.js";
-class u extends l {
-  cFc;
-  cProj;
-  dropout;
-  index;
-  _trainable = !0;
-  constructor(t, e) {
-    super(), this.index = t, this.cFc = i({
-      units: e.mlpFactor * e.nEmbed,
-      activation: "gelu",
-      useBias: e.biasInLinear,
-      kernelInitializer: s({
-        mean: 0,
-        stddev: 0.02
-      }),
-      biasInitializer: "zeros",
-      name: `block_${t}_mlp_cFc`
-    }), this.cProj = i({
-      units: e.nEmbed,
-      useBias: e.biasInLinear,
-      kernelInitializer: s({
-        mean: 0,
-        stddev: 0.02 / Math.sqrt(2 * e.nLayer)
-      }),
-      biasInitializer: "zeros",
-      name: `block_${t}_mlp_cProj`
-    }), this.dropout = c({ rate: e.dropout });
-  }
-  get variables() {
-    return [
-      ...this.cFc.trainableWeights.map((t) => t.read()),
-      ...this.cProj.trainableWeights.map((t) => t.read())
-    ];
-  }
-  get trainable() {
-    return this._trainable;
-  }
-  set trainable(t) {
-    this._trainable = t, this.cFc.trainable = t, this.cProj.trainable = t;
-  }
-  saveWeights(t) {
-    t.set(`block_${this.index}_mlpHidden`, this.cFc.getWeights()), t.set(`block_${this.index}_mlpOut`, this.cProj.getWeights());
-  }
-  loadWeights(t) {
-    this.cFc.setWeights(t.get(`block_${this.index}_mlpHidden`) || []), this.cProj.setWeights(t.get(`block_${this.index}_mlpOut`) || []);
-  }
-  call(t, e = !1) {
-    return n(() => {
-      this.startMemory();
-      const r = this.cFc.apply(t), a = this.cProj.apply(r), o = this.dropout.apply(a, { training: e });
-      return this.endMemory("MLP"), o;
-    });
-  }
-  dispose() {
-    this.cFc.dispose(), this.cProj.dispose(), this.dropout.dispose();
-  }
-}
+import "../index-XjBAhiFO.js";
+import "./BaseLayer.js";
+import "../random_width-CMHmdbSu.js";
+import { M as i } from "../MLP-KHhikThU.js";
 export {
-  u as default
+  i as default
 };

package/dist/layers/RMSNorm.d.ts

@@ -1,9 +1,9 @@
 import { Tensor, Variable } from '@tensorflow/tfjs-core';
-import { default as BaseLayer } from './BaseLayer';
+import { default as BaseLayer, GPTLayerConfig } from './BaseLayer';
 export default class RMSNorm extends BaseLayer {
     private gamma;
     private epsilon;
-    constructor(shape: number[], epsilon?: number, name?: string);
+    constructor(config: GPTLayerConfig, epsilon?: number, name?: string);
     get trainableWeights(): Variable[];
     set trainable(value: boolean);
     getWeights(): Tensor[];

package/dist/layers/RMSNorm.js

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

package/dist/layers/RoPECache.js

@@ -1,6 +1,6 @@
-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";
+import { o as h, h as c, E as f, N as l, f as n, O as m, t as u, F as p } from "../index-XjBAhiFO.js";
+import { c as d, s as C } from "../sin-BYM-U4Ut.js";
+import { r as a } from "../range-DQMNzBWs.js";
 /**
  * @license
  * Copyright 2018 Google LLC. All Rights Reserved.

package/dist/layers/TiedEmbedding.js

@@ -1,8 +1,9 @@
-import { r as t, d as s } from "../random_width-oeUIlUZj.js";
-import "../index-pWA4_lUh.js";
-import { v as r } from "../variable-Dl_ub3pk.js";
-import { g as d } from "../gather-BPGW8RsB.js";
-class b {
+import { r as t } from "../random_width-CMHmdbSu.js";
+import "../index-XjBAhiFO.js";
+import { d as s } from "../tfjs_backend-Bzl2SrRo.js";
+import { v as r } from "../variable-BS4AKqNU.js";
+import { g as d } from "../gather-BUmJIS8n.js";
+class n {
   vocabSize;
   embedDim;
   tiedWeights;
@@ -43,5 +44,5 @@ class b {
   }
 }
 export {
-  b as default
+  n as default
 };

package/dist/layers/TransformerBlock.d.ts

@@ -1,8 +1,5 @@
-import { GPTConfig } from '../config';
 import { KVCache } from './CausalSelfAttention';
-import { default as RoPECache } from './RoPECache';
-import { default as MemoryProfiler } from '../utilities/profile';
-import { default as BaseLayer } from './BaseLayer';
+import { default as BaseLayer, GPTLayerConfig } from './BaseLayer';
 import { Tensor, Variable } from '@tensorflow/tfjs-core';
 export default class Block extends BaseLayer {
     private ln1;
@@ -12,8 +9,7 @@ export default class Block extends BaseLayer {
     private index;
     private _trainable;
     skipped: boolean;
-    constructor(index: number, config: GPTConfig, ropeCache?: RoPECache);
-    setProfiler(value: MemoryProfiler | undefined): void;
+    constructor(index: number, config: GPTLayerConfig);
     get variables(): Variable[];
     get trainable(): boolean;
     set trainable(value: boolean);

package/dist/layers/TransformerBlock.js

@@ -1,9 +1,9 @@
 import h from "./CausalSelfAttention.js";
-import o from "./MLP.js";
-import r from "./RMSNorm.js";
+import { M as o } from "../MLP-KHhikThU.js";
+import a from "./RMSNorm.js";
 import p from "./BaseLayer.js";
-import { t as d } from "../index-pWA4_lUh.js";
-class g extends p {
+import { t as d } from "../index-XjBAhiFO.js";
+class W extends p {
   ln1;
   attn;
   ln2;
@@ -11,11 +11,8 @@ class g extends p {
   index;
   _trainable = !0;
   skipped = !1;
-  constructor(t, s, i) {
-    super(), this.index = t, this.ln1 = new r([s.nEmbed], 1e-8, `block_${this.index}_rms1`), this.attn = new h(this.index, s, i), this.ln2 = new r([s.nEmbed], 1e-8, `block_${this.index}_rms2`), this.mlp = new o(this.index, s);
-  }
-  setProfiler(t) {
-    this._profiler = t, this.attn.setProfiler(t), this.mlp.setProfiler(t), this.ln1.setProfiler(t), this.ln2.setProfiler(t);
+  constructor(t, s) {
+    super(s), this.index = t, this.ln1 = new a(s, 1e-8, `block_${this.index}_rms1`), this.attn = new h(this.index, s), this.ln2 = new a(s, 1e-8, `block_${this.index}_rms2`), this.mlp = new o(this.index, s);
   }
   get variables() {
     return [
@@ -45,9 +42,9 @@ class g extends p {
     return d(() => {
       if (this.skipped)
         return { output: t };
-      const l = this.ln1.apply(t), n = this.attn.call(l, s, i, e), a = t.add(n.output);
+      const l = this.ln1.apply(t), n = this.attn.call(l, s, i, e), r = t.add(n.output);
       return {
-        output: this.getMLPOutput(a, s),
+        output: this.getMLPOutput(r, s),
         attention: n.attention,
         cache: n.presentKV
       };
@@ -58,5 +55,5 @@ class g extends p {
   }
 }
 export {
-  g as default
+  W as default
 };