@genai-fi/nanogpt 0.9.1 → 0.10.1

package/dist/utilities/packed.js

@@ -0,0 +1,716 @@
+import { w as T, o as R, p as x, K as j, I as M, q as V, s as $, t as O, v as L, x as _, A as U } from "../tensor_util-DV-FP5Q3.js";
+import { b as q, o as W, E as H, q as X, a as p, r as Y, t as Z, u as J, v as K, V as Q, w as N, T as B, x as P, f as tt, m as et, s as st, y as nt } from "../tensor-DdQUJZlz.js";
+import { PackableTensor as A, PackableVariable as rt } from "../patches/PackedTensor.js";
+function lt() {
+  return y.backendName === "webgpu";
+}
+function E(o) {
+  return o.packed !== void 0;
+}
+function w(o) {
+  return E(o) && o.packed;
+}
+function z(o) {
+  if (E(o)) {
+    if (o.dtype !== "int32")
+      throw new Error("packTensor: only int32 tensors can be packed.");
+    return o.packed = !0, o;
+  } else
+    throw console.error("Tensor:", o), new Error("Tensor is not packable");
+}
+function ft(o) {
+  if (E(o)) {
+    if (o.dtype !== "float32")
+      throw new Error("unpackTensor: only float32 tensors can be unpacked.");
+    o.packed = !1;
+  }
+  return o;
+}
+function it(o, t, e, s) {
+  for (let n = t.length - 1; n >= 0; n--) {
+    const r = t[n], c = [];
+    if (r.outputs.forEach((a) => {
+      const i = o[a.id];
+      i != null ? c.push(i) : c.push(null);
+    }), r.gradient == null)
+      throw new Error(`Cannot compute gradient: gradient function not found for ${r.kernelName}.`);
+    const h = r.gradient(c);
+    for (const a in r.inputs) {
+      if (!(a in h))
+        throw new Error(
+          `Cannot backprop through input ${a}. Available gradients found: ${Object.keys(h)}.`
+        );
+      const i = e(() => h[a]()), d = w(i);
+      if (i.dtype !== "float32" && (!d || i.dtype !== "int32"))
+        throw new Error(
+          `Error in gradient for op ${r.kernelName}. The gradient of input ${a} must have 'float32' dtype, but has '${i.dtype}'`
+        );
+      const l = r.inputs[a];
+      if (!q(i.shape, l.shape))
+        throw new Error(
+          `Error in gradient for op ${r.kernelName}. The gradient of input '${a}' has shape '${i.shape}', which does not match the shape of the input '${l.shape}'`
+        );
+      if (o[l.id] == null)
+        o[l.id] = i;
+      else {
+        const u = o[l.id];
+        o[l.id] = s(u, i), u.dispose();
+      }
+    }
+  }
+}
+function S(o) {
+  return o.kernelName != null;
+}
+class C {
+  // Public since optimizers will use it.
+  registeredVariables = {};
+  nextTapeNodeId = 0;
+  numBytes = 0;
+  numTensors = 0;
+  numStringTensors = 0;
+  numDataBuffers = 0;
+  activeTape;
+  // Number of nested tf.grad() statements when computing higher-order
+  // gradients. E.g. `1` for first-order gradients and `2` for second-order
+  // gradients. Used to track if the tape should be removed after a backprop.
+  gradientDepth = 0;
+  // Number of nested kernel calls. When kernel depth is greater than 1, we turn
+  // off the tape.
+  kernelDepth = 0;
+  // Keep Tensors that parallel the tapes.
+  activeScope;
+  scopeStack = [];
+  /**
+   * Keeps track of the number of data moves during a kernel execution. We
+   * maintain a stack since kernels can call other kernels, recursively.
+   */
+  numDataMovesStack = [];
+  nextScopeId = 0;
+  tensorInfo = /* @__PURE__ */ new WeakMap();
+  profiling = !1;
+  activeProfile = {
+    newBytes: 0,
+    newTensors: 0,
+    peakBytes: 0,
+    kernels: [],
+    result: null,
+    get kernelNames() {
+      return Array.from(new Set(this.kernels.map((t) => t.name)));
+    }
+  };
+  dispose() {
+    for (const t in this.registeredVariables)
+      this.registeredVariables[t].dispose();
+  }
+}
+class v {
+  constructor(t) {
+    this.ENV = t, this.state = new C(), console.log("GenAI Patched Engine Initialized");
+  }
+  version = "GENAI_PATCHED_ENGINE";
+  state;
+  backendName;
+  registry = {};
+  registryFactory = {};
+  profiler;
+  backendInstance = null;
+  pendingBackendInit;
+  pendingBackendInitId = 0;
+  async ready() {
+    if (this.pendingBackendInit != null)
+      return this.pendingBackendInit.then(() => {
+      });
+    if (this.backendInstance != null)
+      return;
+    const t = this.getSortedBackends();
+    for (let e = 0; e < t.length; e++) {
+      const s = t[e];
+      if (await this.initializeBackend(s).success) {
+        await this.setBackend(s);
+        return;
+      }
+    }
+    throw new Error("Could not initialize any backends, all backend initializations failed.");
+  }
+  get backend() {
+    if (this.pendingBackendInit != null)
+      throw new Error(
+        `Backend '${this.backendName}' has not yet been initialized. Make sure to await tf.ready() or await tf.setBackend() before calling other methods`
+      );
+    if (this.backendInstance == null) {
+      const { name: t, asyncInit: e } = this.initializeBackendsAndReturnBest();
+      if (e)
+        throw new Error(
+          `The highest priority backend '${t}' has not yet been initialized. Make sure to await tf.ready() or await tf.setBackend() before calling other methods`
+        );
+      this.setBackend(t);
+    }
+    return this.backendInstance;
+  }
+  backendNames() {
+    return Object.keys(this.registryFactory);
+  }
+  findBackend(t) {
+    if (!(t in this.registry))
+      if (t in this.registryFactory) {
+        const { asyncInit: e } = this.initializeBackend(t);
+        if (e)
+          return null;
+      } else
+        return null;
+    return this.registry[t];
+  }
+  findBackendFactory(t) {
+    return t in this.registryFactory ? this.registryFactory[t].factory : null;
+  }
+  registerBackend(t, e, s = 1) {
+    return t in this.registryFactory ? (T(`${t} backend was already registered. Reusing existing backend factory.`), !1) : (this.registryFactory[t] = { factory: e, priority: s }, console.log("Registered backend", t), !0);
+  }
+  async setBackend(t) {
+    if (this.registryFactory[t] == null)
+      throw new Error(`Backend name '${t}' not found in registry`);
+    if (this.backendName = t, this.registry[t] == null) {
+      this.backendInstance = null;
+      const { success: e, asyncInit: s } = this.initializeBackend(t);
+      if (!(s ? await e : e))
+        return !1;
+    }
+    return this.backendInstance = this.registry[t], this.setupRegisteredKernels(), this.profiler = new R(this.backendInstance), !0;
+  }
+  setupRegisteredKernels() {
+    x(this.backendName).forEach((e) => {
+      e.setupFunc != null && e.setupFunc(this.backendInstance);
+    });
+  }
+  disposeRegisteredKernels(t) {
+    x(t).forEach((s) => {
+      s.disposeFunc != null && s.disposeFunc(this.registry[t]);
+    });
+  }
+  /**
+   * Initializes a backend by looking up the backend name in the factory
+   * registry and calling the factory method. Returns a boolean representing
+   * whether the initialization of the backend succeeded. Throws an error if
+   * there is no backend in the factory registry.
+   */
+  initializeBackend(t) {
+    const e = this.registryFactory[t];
+    if (e == null)
+      throw new Error(`Cannot initialize backend ${t}, no registration found.`);
+    try {
+      const s = e.factory();
+      if (s && !(s instanceof j) && typeof s.then == "function") {
+        const n = ++this.pendingBackendInitId, r = s.then((c) => n < this.pendingBackendInitId ? !1 : (this.registry[t] = c, this.pendingBackendInit = null, !0)).catch((c) => (n < this.pendingBackendInitId || (this.pendingBackendInit = null, T(`Initialization of backend ${t} failed`), T(c.stack || c.message)), !1));
+        return this.pendingBackendInit = r, { success: r, asyncInit: !0 };
+      } else
+        return this.registry[t] = s, { success: !0, asyncInit: !1 };
+    } catch (s) {
+      return T(`Initialization of backend ${t} failed`), T(s.stack || s.message), { success: !1, asyncInit: !1 };
+    }
+  }
+  removeBackend(t) {
+    if (!(t in this.registryFactory))
+      throw new Error(`${t} backend not found in registry`);
+    this.backendName === t && this.pendingBackendInit != null && this.pendingBackendInitId++, t in this.registry && (this.disposeRegisteredKernels(t), this.registry[t].dispose(), delete this.registry[t]), delete this.registryFactory[t], this.backendName === t && (this.pendingBackendInit = null, this.backendName = null, this.backendInstance = null);
+  }
+  getSortedBackends() {
+    if (Object.keys(this.registryFactory).length === 0)
+      throw new Error("No backend found in registry.");
+    return Object.keys(this.registryFactory).sort((t, e) => this.registryFactory[e].priority - this.registryFactory[t].priority);
+  }
+  initializeBackendsAndReturnBest() {
+    const t = this.getSortedBackends();
+    for (let e = 0; e < t.length; e++) {
+      const s = t[e], { success: n, asyncInit: r } = this.initializeBackend(s);
+      if (r || n)
+        return { name: s, asyncInit: r };
+    }
+    throw new Error("Could not initialize any backends, all backend initializations failed.");
+  }
+  moveData(t, e) {
+    const s = this.state.tensorInfo.get(e);
+    s || console.warn("Tried to move data that does not exist", this.state, e);
+    const n = s.backend, r = this.readSync(e), c = n.refCount(e);
+    n.disposeData(e, !0), s.backend = t, t.move(e, r, s.shape, s.dtype, c), this.shouldCheckForMemLeaks() && this.state.numDataMovesStack[this.state.numDataMovesStack.length - 1]++;
+  }
+  tidy(t, e) {
+    let s = null;
+    if (e == null) {
+      if (typeof t != "function")
+        throw new Error("Please provide a function to tidy()");
+      e = t;
+    } else {
+      if (typeof t != "string" && !(t instanceof String))
+        throw new Error("When calling with two arguments, the first argument to tidy() must be a string");
+      if (typeof e != "function")
+        throw new Error("When calling with two arguments, the 2nd argument to tidy() must be a function");
+      s = t;
+    }
+    let n;
+    return this.scopedRun(
+      () => this.startScope(s),
+      () => this.endScope(n),
+      () => (n = e(), n instanceof Promise && console.error("Cannot return a Promise inside of tidy."), n)
+    );
+  }
+  scopedRun(t, e, s) {
+    t();
+    try {
+      const n = s();
+      return e(), n;
+    } catch (n) {
+      throw e(), n;
+    }
+  }
+  static nextTensorId = 0;
+  nextTensorId() {
+    return v.nextTensorId++;
+  }
+  static nextVariableId = 0;
+  nextVariableId() {
+    return v.nextVariableId++;
+  }
+  /**
+   * This method is called instead of the public-facing tensor.clone() when
+   * saving a tensor for backwards pass. It makes sure to add the clone
+   * operation to the tape regardless of being called inside a kernel
+   * execution.
+   */
+  clone(t) {
+    const s = w(t) ? z(y.runKernel(M, { x: t })) : y.runKernel(M, { x: t }), n = { x: t }, r = (h) => ({
+      x: () => {
+        const a = "float32", i = { x: h }, d = { dtype: a }, l = w(t), u = y.runKernel(
+          _,
+          i,
+          // tslint:disable-next-line: no-unnecessary-type-assertion
+          d
+        );
+        return l && z(u), u;
+      }
+    }), c = [];
+    return this.addTapeNode(this.state.activeScope.name, n, [s], r, c, {}), s;
+  }
+  /**
+   * Execute a kernel with the given name and return the output tensor.
+   *
+   * @param kernelName The name of the kernel to execute.
+   * @param inputs A map of input names to tensors.
+   * @param attrs A map of attribute names to their values. An attribute is a
+   *     primitive (non-tensor) input to the kernel.
+   * @param inputsToSave A list of tensors, inputs to save for the backprop
+   *     computation.
+   * @param outputsToSave A list of booleans, specifying which output to save
+   *     for the backprop computation. These are booleans since the output
+   * tensors are not visible to the user.
+   */
+  runKernel(t, e, s) {
+    if (this.backendName == null && this.backend, !(V(t, this.backendName) != null))
+      throw new Error(`Kernel '${t}' not registered for backend '${this.backendName}'`);
+    return this.runKernelFunc({ kernelName: t, inputs: e, attrs: s });
+  }
+  shouldCheckForMemLeaks() {
+    return this.ENV.getBool("IS_TEST");
+  }
+  checkKernelForMemLeak(t, e, s) {
+    const n = this.backend.numDataIds();
+    let r = 0;
+    s.forEach((a) => {
+      r += a.dtype === "complex64" ? 3 : 1;
+    });
+    const c = this.state.numDataMovesStack[this.state.numDataMovesStack.length - 1], h = n - e - r - c;
+    if (h > 0)
+      throw new Error(
+        `Backend '${this.backendName}' has an internal memory leak (${h} data ids) after running '${t}'`
+      );
+  }
+  /**
+   * Internal helper method to execute a kernel Func
+   *
+   * Use `runKernel` to execute kernels from outside of engine.
+   */
+  runKernelFunc(t) {
+    let e, s = [];
+    const n = this.isTapeOn(), r = this.state.numBytes, c = this.state.numTensors;
+    this.shouldCheckForMemLeaks() && this.state.numDataMovesStack.push(0);
+    let h;
+    this.backendName == null && this.backend;
+    let a;
+    const i = S(t) ? t.kernelName : this.state.activeScope != null ? this.state.activeScope.name : "";
+    if (S(t)) {
+      const { kernelName: f, inputs: I, attrs: m } = t;
+      this.backendName == null && this.backend;
+      const g = V(f, this.backendName);
+      p(
+        g != null,
+        () => `Cannot find registered kernel '${f}' for backend '${this.backendName}'`
+      ), h = () => {
+        const G = this.backend.numDataIds();
+        a = g.kernelFunc({ inputs: I, attrs: m, backend: this.backend });
+        const F = Array.isArray(a) ? a : [a];
+        this.shouldCheckForMemLeaks() && this.checkKernelForMemLeak(f, G, F);
+        const D = F.map((b) => b.rank != null ? b : this.makeTensorFromTensorInfo(b));
+        if (n) {
+          const b = this.getTensorsForGradient(f, I, D);
+          s = this.saveTensorsForBackwardMode(b ?? []);
+        }
+        return D;
+      };
+    } else {
+      const { forwardFunc: f } = t, I = (m) => {
+        n && (s = m.map((g) => this.keep(this.clone(g))));
+      };
+      h = () => {
+        const m = this.backend.numDataIds();
+        a = this.tidy(() => f(this.backend, I));
+        const g = Array.isArray(a) ? a : [a];
+        return this.shouldCheckForMemLeaks() && this.checkKernelForMemLeak(i, m, g), g;
+      };
+    }
+    const { inputs: d, attrs: l } = t, u = S(t) ? null : t.backwardsFunc;
+    let k;
+    return this.scopedRun(
+      // Stop recording to a tape when running a kernel.
+      () => this.state.kernelDepth++,
+      () => this.state.kernelDepth--,
+      () => {
+        !this.ENV.getBool("DEBUG") && !this.state.profiling ? e = h() : (k = this.profiler.profileKernel(i, d, () => h()), this.ENV.getBool("DEBUG") && this.profiler.logKernelProfile(k), e = k.outputs);
+      }
+    ), n && this.addTapeNode(
+      i,
+      d,
+      e,
+      u,
+      s,
+      l ?? {}
+    ), this.state.profiling && this.state.activeProfile.kernels.push({
+      name: i,
+      bytesAdded: this.state.numBytes - r,
+      totalBytesSnapshot: this.state.numBytes,
+      tensorsAdded: this.state.numTensors - c,
+      totalTensorsSnapshot: this.state.numTensors,
+      inputShapes: Object.keys(d).map(
+        (f) => d[f] != null ? d[f].shape : null
+      ),
+      outputShapes: e.map((f) => f.shape),
+      kernelTimeMs: k.timeMs,
+      extraInfo: k.extraInfo
+    }), Array.isArray(a) ? e : e[0];
+  }
+  /**
+   * Saves tensors used in forward mode for use in backward mode.
+   *
+   * @param tensors the list of tensors to save.
+   */
+  saveTensorsForBackwardMode(t) {
+    return t.map((s) => this.keep(this.clone(s)));
+  }
+  /**
+   * Returns a list of tensors to save for a given gradient calculation.
+   *
+   * @param kernelName name of kernel to look up gradient for.
+   * @param inputs a map of input tensors.
+   * @param outputs an array of output tensors from forward mode of kernel.
+   */
+  getTensorsForGradient(t, e, s) {
+    const n = $(t);
+    if (n != null) {
+      const r = n.inputsToSave || [], c = n.outputsToSave || [];
+      let h;
+      n.saveAllInputs ? (p(Array.isArray(e), () => "saveAllInputs is true, expected inputs to be an array."), h = Object.keys(e).map((i) => e[i])) : h = r.map((i) => e[i]);
+      const a = s.filter((i, d) => c[d]);
+      return h.concat(a);
+    }
+    return [];
+  }
+  /**
+   * Internal method used by public APIs for tensor creation. Makes a new
+   * tensor with the provided shape, dtype and values. It always
+   * creates a new data id and writes the values to the underlying backend.
+   */
+  makeTensor(t, e, s, n) {
+    if (t == null)
+      throw new Error("Values passed to engine.makeTensor() are null");
+    s = s || "float32", n = n || this.backend;
+    let r = t;
+    s === "string" && Y(t[0]) && (r = t.map((a) => Z(a)));
+    const c = n.write(r, e, s), h = new A(e, s, c, this.nextTensorId());
+    if (this.trackTensor(h, n), s === "string") {
+      const a = this.state.tensorInfo.get(c), i = J(r);
+      this.state.numBytes += i - a.bytes, a.bytes = i;
+    }
+    return h;
+  }
+  /**
+   * Internal method used by backends. Makes a new tensor
+   * that is a wrapper around an existing data id. It doesn't create
+   * a new data id, only increments the ref count used in memory tracking.
+   * @deprecated
+   */
+  makeTensorFromDataId(t, e, s, n) {
+    s = s || "float32";
+    const r = { dataId: t, shape: e, dtype: s };
+    return this.makeTensorFromTensorInfo(r, n);
+  }
+  /**
+   * Internal method used by backends. Makes a new tensor that is a wrapper
+   * around an existing data id in TensorInfo. It doesn't create a new data id,
+   * only increments the ref count used in memory tracking.
+   */
+  makeTensorFromTensorInfo(t, e) {
+    const { dataId: s, shape: n, dtype: r } = t, c = new A(n, r, s, this.nextTensorId());
+    if (c.packed = t.packed || !1, c.packed && r !== "int32")
+      throw new Error("Only int32 tensors can be packed.");
+    return this.trackTensor(c, e ?? this.backend), c;
+  }
+  makeVariable(t, e = !0, s, n) {
+    s = s || this.nextVariableId().toString(), n != null && n !== t.dtype && (t = t.cast(n));
+    const r = new rt(t, e, s, this.nextTensorId());
+    if (this.state.registeredVariables[r.name] != null)
+      throw new Error(`Variable with name ${r.name} was already registered`);
+    return this.state.registeredVariables[r.name] = r, this.incRef(r, this.backend), r;
+  }
+  trackTensor(t, e) {
+    this.state.numTensors++, t.dtype === "string" && this.state.numStringTensors++;
+    let s = 0;
+    t.dtype !== "complex64" && t.dtype !== "string" && (s = t.size * K(t.dtype)), this.state.numBytes += s, this.state.tensorInfo.has(t.dataId) || (this.state.numDataBuffers++, this.state.tensorInfo.set(t.dataId, {
+      backend: e || this.backend,
+      dtype: t.dtype,
+      shape: t.shape,
+      bytes: s
+    })), t instanceof Q || this.track(t);
+  }
+  // Track the tensor by dataId and increase the refCount for the dataId in the
+  // backend.
+  // TODO(pyu10055): This is currently used by makeVariable method, to increase
+  // refCount on the backend for the dataId. It can potentially be replaced with
+  // Identity op indead of calling backend directly.
+  incRef(t, e) {
+    this.trackTensor(t, e), this.backend.incRef(t.dataId);
+  }
+  removeDataId(t, e) {
+    this.state.tensorInfo.has(t) && this.state.tensorInfo.get(t).backend === e && (this.state.tensorInfo.delete(t), this.state.numDataBuffers--);
+  }
+  disposeTensor(t) {
+    if (!this.state.tensorInfo.has(t.dataId))
+      return;
+    const e = this.state.tensorInfo.get(t.dataId);
+    if (this.state.numTensors--, t.dtype === "string" && (this.state.numStringTensors--, this.state.numBytes -= e.bytes), t.dtype !== "complex64" && t.dtype !== "string") {
+      const s = t.size * K(t.dtype);
+      this.state.numBytes -= s;
+    }
+    e.backend.disposeData(t.dataId) && this.removeDataId(t.dataId, e.backend);
+  }
+  disposeVariables() {
+    for (const t in this.state.registeredVariables) {
+      const e = this.state.registeredVariables[t];
+      this.disposeVariable(e);
+    }
+  }
+  disposeVariable(t) {
+    this.disposeTensor(t), this.state.registeredVariables[t.name] != null && delete this.state.registeredVariables[t.name];
+  }
+  memory() {
+    const t = this.backend.memory();
+    return t.numTensors = this.state.numTensors, t.numDataBuffers = this.state.numDataBuffers, t.numBytes = this.state.numBytes, this.state.numStringTensors > 0 && (t.unreliable = !0, t.reasons == null && (t.reasons = []), t.reasons.push("Memory usage by string tensors is approximate (2 bytes per character)")), t;
+  }
+  async profile(t) {
+    this.state.profiling = !0;
+    const e = this.state.numBytes, s = this.state.numTensors;
+    this.state.activeProfile.kernels = [], this.state.activeProfile.result = await t(), this.state.profiling = !1, this.state.activeProfile.peakBytes = Math.max(
+      ...this.state.activeProfile.kernels.map((n) => n.totalBytesSnapshot)
+    ), this.state.activeProfile.newBytes = this.state.numBytes - e, this.state.activeProfile.newTensors = this.state.numTensors - s;
+    for (const n of this.state.activeProfile.kernels)
+      n.kernelTimeMs = await n.kernelTimeMs, n.extraInfo = await n.extraInfo;
+    return this.state.activeProfile;
+  }
+  isTapeOn() {
+    return this.state.gradientDepth > 0 && this.state.kernelDepth === 0;
+  }
+  addTapeNode(t, e, s, n, r, c) {
+    const h = { id: this.state.nextTapeNodeId++, kernelName: t, inputs: e, outputs: s, saved: r }, a = $(t);
+    a != null && (n = a.gradFunc), n != null && (h.gradient = (i) => (i = i.map((d, l) => {
+      if (d == null) {
+        const u = s[l], k = tt(u.size, u.dtype);
+        return this.makeTensor(k, u.shape, u.dtype);
+      }
+      return d;
+    }), n(i.length > 1 ? i : i[0], r, c))), this.state.activeTape.push(h);
+  }
+  keep(t) {
+    return t.kept = !0, t;
+  }
+  startTape() {
+    this.state.gradientDepth === 0 && (this.state.activeTape = []), this.state.gradientDepth++;
+  }
+  endTape() {
+    this.state.gradientDepth--;
+  }
+  /**
+   * Start a scope. Use this with endScope() to achieve the same functionality
+   * as scope() without the need for a function closure.
+   */
+  startScope(t) {
+    const e = {
+      track: [],
+      name: "unnamed scope",
+      id: this.state.nextScopeId++
+    };
+    t && (e.name = t), this.state.scopeStack.push(e), this.state.activeScope = e;
+  }
+  /**
+   * End a scope. Use this with startScope() to achieve the same functionality
+   * as scope() without the need for a function closure.
+   */
+  endScope(t) {
+    const e = O(t), s = new Set(e.map((r) => r.id));
+    for (let r = 0; r < this.state.activeScope.track.length; r++) {
+      const c = this.state.activeScope.track[r];
+      !c.kept && !s.has(c.id) && c.dispose();
+    }
+    const n = this.state.scopeStack.pop();
+    this.state.activeScope = this.state.scopeStack.length === 0 ? null : this.state.scopeStack[this.state.scopeStack.length - 1], e.forEach((r) => {
+      !r.kept && r.scopeId === n?.id && this.track(r);
+    });
+  }
+  /**
+   * Returns gradients of `f` with respect to each of the `xs`. The gradients
+   * returned are of the same length as `xs`, but some might be null if `f`
+   * was not a function of that `x`. It also takes optional dy to multiply the
+   * gradient, which defaults to `1`.
+   */
+  gradients(t, e, s, n = !1) {
+    if (p(e.length > 0, () => "gradients() received an empty list of xs."), s != null && s.dtype !== "float32")
+      throw new Error(`dy must have 'float32' dtype, but has '${s.dtype}'`);
+    const r = this.scopedRun(
+      () => this.startTape(),
+      () => this.endTape(),
+      () => this.tidy("forward", t)
+    );
+    p(r instanceof B, () => "The result y returned by f() must be a tensor.");
+    const c = L(this.state.activeTape, e, r);
+    if (!n && c.length === 0 && e.length > 0)
+      throw new Error(
+        "Cannot compute gradient of y=f(x) with respect to x. Make sure that the f you passed encloses all operations that lead from x to y."
+      );
+    return this.tidy("backward", () => {
+      const h = {};
+      h[r.id] = s ?? at(r.shape), it(
+        h,
+        c,
+        // Pass the tidy function to avoid circular dep with `tape.ts`.
+        (i) => this.tidy(i),
+        // Pass an add function to avoide a circular dep with `tape.ts`.
+        ct
+      );
+      const a = e.map((i) => h[i.id]);
+      return this.state.gradientDepth === 0 && (this.state.activeTape.forEach((i) => {
+        if (i.saved !== void 0)
+          for (const d of i.saved)
+            d.dispose();
+      }), this.state.activeTape = null), { value: r, grads: a };
+    });
+  }
+  customGrad(t) {
+    return p(N(t), () => "The f passed in customGrad(f) must be a function."), (...e) => {
+      p(
+        e.every((h) => h instanceof B),
+        () => "The args passed in customGrad(f)(x1, x2,...) must all be tensors"
+      );
+      let s;
+      const n = {};
+      e.forEach((h, a) => {
+        n[a] = h;
+      });
+      const r = (h, a) => (s = t(...e, a), p(
+        s.value instanceof B,
+        () => "The function f passed in customGrad(f) must return an object where `obj.value` is a tensor"
+      ), p(
+        N(s.gradFunc),
+        () => "The function f passed in customGrad(f) must return an object where `obj.gradFunc` is a function."
+      ), s.value), c = (h, a) => {
+        const i = s.gradFunc(h, a), d = Array.isArray(i) ? i : [i];
+        p(
+          d.length === e.length,
+          () => "The function f passed in customGrad(f) must return an object where `obj.gradFunc` is a function that returns the same number of tensors as inputs passed to f(...)."
+        ), p(
+          d.every((u) => u instanceof B),
+          () => "The function f passed in customGrad(f) must return an object where `obj.gradFunc` is a function that returns a list of only tensors."
+        );
+        const l = {};
+        return d.forEach((u, k) => {
+          l[k] = () => u;
+        }), l;
+      };
+      return this.runKernelFunc({
+        forwardFunc: r,
+        backwardsFunc: c,
+        inputs: n
+      });
+    };
+  }
+  readSync(t) {
+    return this.state.tensorInfo.get(t).backend.readSync(t);
+  }
+  read(t) {
+    return this.state.tensorInfo.get(t).backend.read(t);
+  }
+  readToGPU(t, e) {
+    return this.state.tensorInfo.get(t).backend.readToGPU(t, e);
+  }
+  async time(t) {
+    const e = P(), s = await this.backend.time(t);
+    return s.wallMs = P() - e, s;
+  }
+  /**
+   * Tracks a Tensor in the current scope to be automatically cleaned up
+   * when the current scope ends, and returns the value.
+   *
+   * @param result The Tensor to track in the current scope.
+   */
+  track(t) {
+    return this.state.activeScope != null && (t.scopeId = this.state.activeScope.id, this.state.activeScope.track.push(t)), t;
+  }
+  get registeredVariables() {
+    return this.state.registeredVariables;
+  }
+  /**
+   * Resets the engine state. Removes all backends but does not remove
+   * registered backend factories.
+   */
+  reset() {
+    this.pendingBackendInitId++, this.state.dispose(), this.ENV.reset(), this.state = new C();
+    for (const t in this.registry)
+      this.disposeRegisteredKernels(t), this.registry[t].dispose(), delete this.registry[t];
+    this.backendName = null, this.backendInstance = null, this.pendingBackendInit = null;
+  }
+}
+function at(o) {
+  const t = et(st(o), "float32");
+  return y.makeTensor(t, o, "float32");
+}
+function ot() {
+  const o = W();
+  if (o._tfengine == null) {
+    const t = new H(o);
+    o._tfengine = new v(t);
+  }
+  return X(o._tfengine.ENV), nt(() => o._tfengine), o._tfengine;
+}
+const y = ot();
+function ct(o, t) {
+  const e = w(o) || w(t), s = { a: o, b: t };
+  return y.runKernel(e ? "Add16" : U, s);
+}
+export {
+  v as E,
+  y as a,
+  it as b,
+  ct as c,
+  ot as g,
+  E as isPackableTensor,
+  w as isPackedTensor,
+  z as packTensor,
+  lt as packingSupported,
+  ft as unpackTensor
+};