import { b as m, t as n } from "../index-DSGwv2Yx.js";
import { p } from "../index-DmeWGGmS.js";
import { v as g } from "../variable-SSATClyt.js";
import { r as S } from "../random_normal-CIm8lk2-.js";
import { z as _ } from "../zeros-Bw0puq_w.js";
class B {
  weightStore;
  alpha;
  rank;
  variables;
  scale;
  name;
  constructor(t, e, a, s, r) {
    this.name = t, this.weightStore = e, this.alpha = a, this.rank = s;
    const c = p(r), w = e.variableNames.filter(
      (i) => c(i) && !i.endsWith("_loraA") && !i.endsWith("_loraB")
    );
    this.variables = new Set(w), this.scale = m(a / s), this.variables.forEach((i) => {
      const o = this.weightStore.getRawVariable(i), [d, b] = o.shape, h = `${i}_${this.name}_loraA`, l = `${i}_${this.name}_loraB`;
      if (o.shape.length !== 2) {
        console.warn(
          `LoRA currently only supports 2D weight matrices. Variable ${i} has shape ${o.shape}`
        ), this.variables.delete(i);
        return;
      }
      this.weightStore.hasVariable(h) || this.weightStore.hasVariable(l) || (this.weightStore.addVariable(
        h,
        g(S([d, this.rank], 0, 0.02), !0, h)
      ), this.weightStore.addVariable(l, g(_([this.rank, b]), !0, l)));
    });
  }
  attach() {
    if (this.weightStore.onWeightRead)
      throw new Error("LoRA cannot be applied to a WeightStore that already has a onWeightRead hook.");
    this.weightStore.onWeightRead = (t, e) => this.variables.has(t) ? n(() => {
      const a = this.weightStore.getRawVariable(`${t}_${this.name}_loraA`), s = this.weightStore.getRawVariable(`${t}_${this.name}_loraB`);
      return e.add(a.matMul(s).mul(this.scale));
    }) : e, this.weightStore.setTrainable([`*_${this.name}_loraA`, `*_${this.name}_loraB`]);
  }
  merge() {
    this.variables.forEach((t) => {
      const e = this.weightStore.getRawVariable(t), a = this.weightStore.getRawVariable(`${t}_${this.name}_loraA`), s = this.weightStore.getRawVariable(`${t}_${this.name}_loraB`), r = n(() => e.add(a.matMul(s).mul(this.scale)));
      e.assign(r), r.dispose();
    });
  }
  detach() {
    this.weightStore.onWeightRead = void 0, this.weightStore.setTrainable(["*"]);
  }
  dispose() {
    this.detach(), this.scale.dispose(), this.variables.forEach((t) => {
      const e = `${t}_${this.name}_loraA`, a = `${t}_${this.name}_loraB`;
      this.weightStore.getRawVariable(e).dispose(), this.weightStore.getRawVariable(a).dispose(), this.weightStore.deleteVariable(e), this.weightStore.deleteVariable(a);
    }), this.variables.clear();
  }
}
export {
  B as default
};
