import { adamAdjust as B } from "../ops/adamAdjust.js";
import { adamMoments as N } from "../ops/adamMoments.js";
import { O as S, e as b, t as c, b as M, l as w } from "../index-DSGwv2Yx.js";
import R from "./LRScheduler.js";
import { clipScale as f } from "../ops/globalNorm.js";
import { save_safetensors as v, load_safetensors as A } from "../utilities/safetensors.js";
import { z as O } from "../zeros-Bw0puq_w.js";
class _ extends S {
  constructor(t) {
    super(), this.config = t, this.accBeta1 = t.accBeta1 ?? t.beta1, this.accBeta2 = t.accBeta2 ?? t.beta2, this.learningRate = t.learningRate, this.beta1 = t.beta1, this.beta2 = t.beta2, this.weightDecay = t.weightDecay, this.lossScaling = t.lossScaling, this.clipNorm = t.clipNorm, this.orthGrad = t.orthoGrad ?? !1, t.epsilon === null || t.epsilon === void 0 ? this.epsilon = b().backend.epsilon() : this.epsilon = t.epsilon, this.lrScheduler = new R(t.learningRate, t);
  }
  className = "AdamW";
  accBeta1 = 0;
  accBeta2 = 0;
  accumulatedMoments = [];
  learningRate;
  beta1;
  beta2;
  lossScaling;
  weightDecay;
  epsilon = null;
  lrScheduler;
  clipNorm;
  orthGradEpsilon = 1e-30;
  orthGrad;
  get lr() {
    return this.learningRate;
  }
  saveMoments() {
    const t = {};
    return this.accumulatedMoments.forEach((e) => {
      t[e.originalName] = e.variable;
    }), v(t);
  }
  async loadMoments(t) {
    const e = await A(t);
    Object.entries(e).forEach(([a, s]) => {
      const n = s.variable(!1);
      this.accumulatedMoments.push({ originalName: a, variable: n });
    });
  }
  serializeConfig() {
    return {
      learningRate: this.learningRate,
      beta1: this.beta1,
      beta2: this.beta2,
      accBeta1: this.accBeta1,
      accBeta2: this.accBeta2,
      epsilon: this.epsilon ?? void 0,
      weightDecay: this.weightDecay,
      lossScaling: this.lossScaling,
      clipNorm: this.clipNorm,
      orthoGrad: this.orthGrad,
      ...this.lrScheduler.serializeConfig()
    };
  }
  orthogonalizeGradient(t, e) {
    return c(() => {
      const a = t.reshape([-1]), s = e.reshape([-1]), n = a.mul(a).sum().add(this.orthGradEpsilon), h = a.mul(s).sum().div(n), o = s.sub(a.mul(h)), l = s.norm(), i = o.norm().add(this.orthGradEpsilon);
      return o.mul(l.div(i)).reshape(e.shape);
    });
  }
  updateConfig(t) {
    const e = { ...this.config, ...t };
    this.learningRate = e.learningRate, this.beta1 = e.beta1, this.beta2 = e.beta2, this.weightDecay = e.weightDecay, this.lossScaling = e.lossScaling, this.epsilon = e.epsilon ?? this.epsilon, this.clipNorm = e.clipNorm, this.lrScheduler.updateConfig(e, e.learningRate);
  }
  applyGradients(t) {
    const e = this.lrScheduler.getNextLR();
    this.learningRate = e;
    const a = Array.isArray(t) ? t.map((n) => n.name) : Object.keys(t), s = c(() => {
      const n = 1 - this.accBeta1, h = 1 - this.accBeta2;
      let o;
      if (this.clipNorm !== void 0) {
        const l = a.map((i, r) => Array.isArray(t) ? t[r].tensor : t[i]);
        o = f(l, 1 / this.lossScaling, this.clipNorm);
      } else
        o = M(1 / this.lossScaling);
      return a.forEach((l, i) => {
        const r = b().registeredVariables[l], p = !1;
        this.accumulatedMoments[i] == null && (this.accumulatedMoments[i] = {
          originalName: `${l}/m`,
          variable: c(() => O([...r.shape, 2]).variable(p))
        });
        const m = Array.isArray(t) ? t[i].tensor : t[l];
        if (m == null)
          return;
        const u = this.orthGrad ? this.orthogonalizeGradient(r, m) : m, d = this.accumulatedMoments[i].variable, g = N(d, u, this.beta1, this.beta2, o);
        d.assign(g), this.orthGrad && u.dispose();
        const y = B(
          g,
          r,
          n,
          h,
          this.epsilon ?? 1e-8,
          this.learningRate,
          // Only apply weight decay if the variable is multi-dimensional (e.g. weights, not biases)
          r.shape.length > 1 ? this.weightDecay : 0
        );
        r.assign(y);
      }), this.accBeta1 = this.accBeta1 * this.beta1, this.accBeta2 = this.accBeta2 * this.beta2, o;
    });
    return this.incrementIterations(), s;
  }
  dispose() {
    this.accumulatedMoments != null && w(this.accumulatedMoments.map((t) => t.variable));
  }
  async getWeights() {
    const t = [...this.accumulatedMoments];
    return [await this.saveIterations()].concat(
      t.map((e) => ({ name: e.originalName, tensor: e.variable }))
    );
  }
  async setWeights(t) {
    t = await this.extractIterations(t), c(() => {
      this.accBeta1 = Math.pow(this.beta1, this.iterations_ + 1), this.accBeta2 = Math.pow(this.beta2, this.iterations_ + 1);
    });
    const e = t.length / 2, a = !1;
    this.accumulatedMoments = t.slice(0, e).map((s) => ({
      originalName: s.name,
      variable: s.tensor.variable(a)
    }));
  }
  getConfig() {
    return {
      learningRate: this.learningRate,
      beta1: this.beta1,
      beta2: this.beta2,
      epsilon: this.epsilon
    };
  }
  /** @nocollapse */
  static fromConfig(t, e) {
    return new t(e.learningRate, e.beta1, e.beta2, e.epsilon);
  }
}
export {
  _ as AdamWOptimizer
};
