import { a2 as p, e as s, a4 as g } from "../index-DSGwv2Yx.js";
import b from "./WeightStore.js";
class T {
  parent;
  config;
  weightStore;
  children = [];
  profiler;
  ownVariables = /* @__PURE__ */ new Set();
  constructor(t, e) {
    this.config = t, this.parent = e, this.parent ? (this.parent.children.push(this), this.weightStore = this.parent.weightStore) : this.weightStore = new b();
  }
  getProfiler() {
    return this.profiler;
  }
  setProfiler(t) {
    this.profiler = t || void 0, this.children.forEach((e) => {
      e.setProfiler(t);
    });
  }
  startMemory() {
    this.profiler?.startMemory();
  }
  endMemory(t) {
    this.profiler?.endMemory(t);
  }
  addVariable(t, e) {
    this.weightStore.addVariable(t, e), this.ownVariables.add(t), this.parent && this.parent.addChildVariable(t);
  }
  addChildVariable(t) {
    this.ownVariables.add(t);
  }
  get variables() {
    return this.weightStore.variables;
  }
  get trainableVariables() {
    return this.weightStore.trainableVariables.filter((t) => this.ownVariables.has(t.name));
  }
  getVariable(t) {
    return this.weightStore.getVariable(t);
  }
  hasVariable(t) {
    return this.weightStore.hasVariable(t);
  }
  setVariable(t, e) {
    this.weightStore.setVariable(t, e);
  }
  dispose() {
    this.weightStore.dispose();
  }
  build() {
  }
  call(t, ...e) {
    return this.build(), this.forward(t, ...e);
  }
  callCheckpoint(t, ...e) {
    return this.build(), this.checkpointingFn(t, ...e);
  }
  checkpointingFn(t, ...e) {
    const r = this.trainableVariables;
    return p((...i) => {
      const o = i[i.length - 1], a = i.slice(0, e.length), n = this.forward(t, ...a);
      return o(a), { value: n, gradFunc: (h, l) => {
        const c = s().state.activeTape;
        s().state.activeTape = [];
        const d = g((...u) => this.forward(t, ...u.slice(0, a.length)))([...l, ...r], h);
        return s().state.activeTape = c, d;
      } };
    })(...e, ...r);
  }
}
export {
  T as default
};
