import { d as a, b as n, p as c, t as p, k as r } from "../index-DSGwv2Yx.js";
import { r as l } from "../zeros_like-rOHr54NY.js";
import { c as f, s as m } from "../unsorted_segment_sum-BfFVV9Zm.js";
import { r as h } from "../range-BDxO73mk.js";
class x {
  constructor(s) {
    this.config = s;
    const o = this.config.nEmbed / this.config.nHead;
    if (this.rotaryDim = o, this.rotaryDim % 2 !== 0)
      throw new Error("rotaryDim must be even");
    this.ropeBase = 1e4;
    const i = h(0, this.rotaryDim, 2, "float32"), e = a(i, n(this.rotaryDim, "float32")), t = c(n(this.ropeBase, "float32"), e);
    this.ropeInvFreq = l(t), e.dispose(), t.dispose(), i.dispose(), (this.config.modelType === "GenAI_NanoGPT_v1" ? this.config.useRope ?? !0 : !0) === !1 ? (this.ropeCos = null, this.ropeSin = null, this.ropeNegSin = null, this.ropeCacheLen = 0) : p(() => {
      this.ensureRopeCache(this.config.blockSize * 4);
    });
  }
  rotaryDim;
  ropeBase;
  ropeInvFreq;
  ropeCos = null;
  // [cacheLen, rotaryDim/2]
  ropeSin = null;
  // [cacheLen, rotaryDim/2]
  ropeNegSin = null;
  // [cacheLen, rotaryDim/2]
  ropeCacheLen = 0;
  ensureRopeCache(s) {
    p(() => {
      if (s <= this.ropeCacheLen) return;
      this.ropeCos && this.ropeCos.dispose(), this.ropeSin && this.ropeSin.dispose();
      const o = Math.max(s, this.ropeCacheLen + this.config.blockSize * 4), e = h(0, o, 1, "float32").expandDims(1).mul(this.ropeInvFreq.expandDims(0));
      this.ropeCos = r(f(e).expandDims(-1)), this.ropeSin = r(m(e).expandDims(-1)), this.ropeNegSin = r(this.ropeSin.neg()), this.ropeCacheLen = o;
    });
  }
  getCos() {
    return this.ropeCos;
  }
  getSin() {
    return this.ropeSin;
  }
  getNegSin() {
    return this.ropeNegSin;
  }
  dispose() {
    this.ropeCos && this.ropeCos.dispose(), this.ropeSin && this.ropeSin.dispose(), this.ropeInvFreq.dispose();
  }
}
export {
  x as default
};
