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 w from "../layers/PositionEmbedding.js";
import { packingSupported as E } from "../utilities/packed.js";
import { p as _, u as C } from "../pack16-WlOSOuZA.js";
const P = {
  modelType: "GenAI_NanoGPT_v1",
  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,
  useRope: !0
};
class S extends k {
  wte;
  // Token embeddings
  wpe;
  // Position embeddings
  // private drop: layers.Layer; // Dropout
  blocks;
  lnF;
  // Final layer norm
  ropeCache;
  constructor(e = {}) {
    super({ ...P, ...e });
    const t = {
      activation: "gelu",
      hiddenFactor: this.config.mlpFactor,
      useGamma: !0,
      useQKNorm: !1
    };
    this.wte = new f(this.config, "token_embedding", this), this.config.useRope === !1 ? this.wpe = new w(this.config, "positional_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, t, this));
    this.lnF = new b(this.config, t, "final_rms_norm", this);
  }
  getClassName() {
    return "GenAI_NanoGPT_v1";
  }
  inputPhase(e, t) {
    return r(() => {
      const o = this.wte.embed(e);
      if (this.config.useRope === !1) {
        const i = this.wpe.call(t, o);
        if (Array.isArray(i))
          throw new Error("PositionEmbedding output should not be an array");
        return i;
      }
      return o;
    });
  }
  forward(e, t) {
    return this.validateInput(t), e.ropeCache = this.ropeCache, e.outputEmbeddings && (e.embeddings = []), r(() => {
      this.startMemory();
      let o = this.inputPhase(t, e);
      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 i = e.mixedPrecision === !0 && E();
      let s = i ? _(o) : o;
      i && o !== s && o.dispose();
      for (let n = 0; n < this.blocks.length; n++) {
        const p = this.blocks[n], a = Math.random() * 1e9, l = {
          ...e,
          seed: a,
          pastKV: e.cache ? e.cache[n] : void 0,
          mixedPrecision: i
        }, m = e.checkpointing && e.training ? p.callCheckpoint(l, s) : p.call(l, s);
        e.outputEmbeddings ? (d(s), e.embeddings.push({ name: `block_output_${n}`, tensor: s })) : s.dispose(), s = m;
      }
      if (o = this.lnF.call({ ...e, mixedPrecision: i }, s), s.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 = i ? C(c) : c;
      return i && 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
};
