import u from "../layers/TransformerBlock.js";
import f from "../layers/TiedEmbedding.js";
import g from "../layers/RoPECache.js";
import b from "../layers/RMSNorm.js";
import { t as r, k as d } from "../index-DSGwv2Yx.js";
import k from "./model.js";
import { packingSupported as w } from "../utilities/packed.js";
import { p as C, u as _ } from "../pack16-WlOSOuZA.js";
const E = {
  modelType: "GenAI_NanoGPT_v2",
  vocabSize: 2e3,
  blockSize: 128,
  // Maximum sequence length
  nLayer: 6,
  // Number of transformer layers
  nHead: 4,
  // Number of attention heads
  nEmbed: 256,
  // Embedding dimension
  mlpFactor: 4
};
class S extends k {
  wte;
  // Token embeddings
  wpe;
  // Position embeddings
  // private drop: layers.Layer; // Dropout
  blocks;
  lnF;
  // Final layer norm
  ropeCache;
  constructor(e = {}) {
    super({ ...E, ...e });
    const s = {
      activation: "relu2",
      hiddenFactor: this.config.mlpFactor,
      useGamma: !1,
      useQKNorm: !0
    };
    this.wte = new f(this.config, "token_embedding", this), this.ropeCache = new g(this.config), this.blocks = [];
    for (let o = 0; o < this.config.nLayer; o++)
      this.blocks.push(new u(o, this.config, s, this));
    this.lnF = new b(this.config, s, "final_rms_norm", this);
  }
  getClassName() {
    return "GenAI_NanoGPT_v2";
  }
  inputPhase(e) {
    return r(() => this.wte.embed(e));
  }
  forward(e, s) {
    return this.validateInput(s), e.ropeCache = this.ropeCache, e.outputEmbeddings && (e.embeddings = []), r(() => {
      this.startMemory();
      let o = this.inputPhase(s);
      if (e.cache && e.cache.length !== this.blocks.length)
        throw console.error("Cache", e.cache), new Error(
          `Cache length ${e.cache.length} does not match number of blocks ${this.blocks.length}`
        );
      const t = e.mixedPrecision === !0 && w();
      let i = t ? C(o) : o;
      t && o !== i && o.dispose();
      for (let n = 0; n < this.blocks.length; n++) {
        if (e.layerDrop && Math.random() < e.layerDrop * (n / this.blocks.length))
          continue;
        const p = this.blocks[n], m = Math.random() * 1e9, l = {
          ...e,
          seed: m,
          pastKV: e.cache ? e.cache[n] : void 0,
          mixedPrecision: t
        }, a = e.checkpointing && e.training ? p.callCheckpoint(l, i) : p.call(l, i);
        e.outputEmbeddings ? (d(i), e.embeddings.push({ name: `block_output_${n}`, tensor: i })) : i.dispose(), i = a;
      }
      if (o = this.lnF.call({ ...e, mixedPrecision: t }, i), i.dispose(), e.skipLogits)
        return this.endMemory("Forward"), o;
      const c = this.wte.project(o);
      e.outputEmbeddings ? (d(o), e.embeddings.push({ name: "final_norm_output", tensor: o })) : o.dispose();
      const h = t ? _(c) : c;
      return t && c !== h && c.dispose(), h;
    });
  }
  project(e) {
    return r(() => this.wte.project(e));
  }
  dispose() {
    this.weightStore.dispose(), this.wte.dispose(), this.wpe && this.wpe.dispose(), this.blocks.forEach((e) => e.dispose()), this.lnF.dispose();
  }
}
export {
  S as default
};
