import { u as o } from "../index-DSGwv2Yx.js";
import { p as h } from "../index-DmeWGGmS.js";
import { v as b } from "../variable-SSATClyt.js";
class d {
  _variables = /* @__PURE__ */ new Map();
  touchedVariables = /* @__PURE__ */ new Set();
  // Hooks
  onWeightRead;
  saveWeights(a) {
    this._variables.forEach((e, i) => {
      e && this.touchedVariables.has(i) && a.set(i, [o(e)]);
    });
  }
  loadWeights(a, e, i = !0) {
    a.forEach((s, r) => {
      const t = s[0], l = this._variables.get(r);
      l ? l.assign(t) : this._variables.set(r, b(t, i, r)), e ? this.touchedVariables.delete(r) : this.touchedVariables.add(r);
    });
  }
  addVariable(a, e) {
    this._variables.set(a, e || null);
  }
  deleteVariable(a) {
    const e = this._variables.get(a);
    e && e.dispose(), this._variables.delete(a), this.touchedVariables.delete(a);
  }
  get variables() {
    return Array.from(this._variables.values()).filter((e) => e !== null);
  }
  get variableNames() {
    return Array.from(this._variables.keys());
  }
  get trainableVariables() {
    return Array.from(this._variables.values()).filter(
      (e) => e !== null && e.trainable
    );
  }
  getRawVariable(a) {
    const e = this._variables.get(a);
    if (!e)
      throw new Error(`Variable ${a} not found`);
    return e;
  }
  getVariable(a) {
    const e = this._variables.get(a);
    if (!e)
      throw new Error(`Variable ${a} not found`);
    return this.onWeightRead ? this.onWeightRead(a, e) : e;
  }
  setTrainable(a) {
    const e = h(a);
    this._variables.forEach((i, s) => {
      i && (i.trainable = e(s));
    });
  }
  hasVariable(a) {
    return !!this._variables.get(a);
  }
  setVariable(a, e) {
    if (!this._variables.has(a))
      throw new Error(`Variable ${a} not found`);
    this._variables.set(a, e);
  }
  touchVariables(a) {
    for (const e of a)
      this._variables.get(e) && this.touchedVariables.add(e);
  }
  dispose() {
    this._variables.forEach((a) => {
      a?.dispose();
    }), this._variables.clear();
  }
}
export {
  d as default
};
