import p from "./CausalSelfAttention.js";
import h from "./MLP.js";
import l from "./RMSNorm.js";
import m from "./BaseLayer.js";
import { k as n, t as u } from "../index-DSGwv2Yx.js";
import { add16 as d } from "../ops/add16.js";
class _ extends m {
  ln1;
  attn;
  ln2;
  mlp;
  index;
  skipped = !1;
  blockConfig;
  constructor(i, s, e, t) {
    super(s, t), this.index = i, this.blockConfig = e, this.ln1 = new l(s, this.blockConfig, `block_${this.index}_rms1`, this), this.attn = new p(this.index, s, this.blockConfig, this), this.ln2 = new l(s, this.blockConfig, `block_${this.index}_rms2`, this), this.mlp = new h(this.index, s, this.blockConfig, this);
  }
  getMLPOutput(i, s) {
    const e = this.ln2.call({ training: s.training }, i), t = this.mlp.call(s, e);
    s.outputEmbeddings ? (n(e), s.embeddings.push({ name: `block_ln2_${this.index}`, tensor: e })) : e.dispose();
    const o = d(i, t);
    return i.dispose(), s.outputEmbeddings ? (n(t), s.embeddings.push({ name: `block_mlp_out_${this.index}`, tensor: t })) : t.dispose(), o;
  }
  forward(i, s) {
    return u(() => {
      if (this.skipped)
        return s;
      const e = this.ln1.call(i, s), t = this.attn.call(i, e);
      i.outputEmbeddings ? (n(e), i.embeddings.push({ name: `block_ln1_${this.index}`, tensor: e })) : e.dispose();
      const o = d(s, t);
      return i.outputEmbeddings ? (n(t), i.embeddings.push({ name: `block_attn_out_${this.index}`, tensor: t })) : t.dispose(), this.getMLPOutput(o, i);
    });
  }
  dispose() {
    this.ln1.dispose(), this.attn.dispose(), this.ln2.dispose(), this.mlp.dispose();
  }
}
export {
  _ as default
};
