import { t as M } from "../index-DSGwv2Yx.js";
import f from "./BaseLayer.js";
import { b as h } from "../matMul16-BIT70Vya.js";
import { reshape16 as d } from "../ops/reshape16.js";
import { dropout16 as L } from "../ops/dropout16.js";
import { v as n } from "../variable-SSATClyt.js";
import { r as m } from "../random_normal-CIm8lk2-.js";
class N extends f {
  index;
  hiddenUnits;
  MLPHIDDEN;
  MLPOUT;
  mlpConfig;
  constructor(i, t, s, e) {
    super(t, e), this.index = i, this.mlpConfig = s, this.hiddenUnits = (s.hiddenFactor ?? t.mlpFactor) * t.nEmbed, this.MLPHIDDEN = `block_${this.index}_mlpHidden`, this.MLPOUT = `block_${this.index}_mlpOut`, this.addVariable(this.MLPHIDDEN), this.addVariable(this.MLPOUT);
  }
  build() {
    this.hasVariable(this.MLPHIDDEN) === !1 && this.setVariable(
      this.MLPHIDDEN,
      n(m([this.config.nEmbed, this.hiddenUnits], 0, 0.02), !0, this.MLPHIDDEN)
    ), this.hasVariable(this.MLPOUT) === !1 && this.setVariable(
      this.MLPOUT,
      n(
        m([this.hiddenUnits, this.config.nEmbed], 0, 0.02 / Math.sqrt(2 * this.config.nLayer)),
        !0,
        this.MLPOUT
      )
    );
  }
  forward(i, t) {
    return M(() => {
      this.startMemory();
      const [s, e, r] = t.shape, l = d(t, [s * e, r]), a = h(l, this.getVariable(this.MLPHIDDEN), !1, !1, {
        activation: this.mlpConfig.activation ?? "gelu"
      }), p = h(a, this.getVariable(this.MLPOUT));
      a.dispose();
      const o = d(p, [s, e, r]);
      return this.endMemory("MLP"), i.dropout && i.dropout > 0 ? L(o, i.dropout) : o;
    });
  }
}
export {
  N as default
};
