@genai-fi/nanogpt 0.6.0 → 0.6.2

package/dist/RealDiv-BYViZwhN.js

@@ -0,0 +1,540 @@
+import { ao as T, ac as E, p as O, g as V, aw as B, N as F, M as j, ax as K } from "./index-BAzbokzv.js";
+import { r as $ } from "./Reshape-t7Kcikjk.js";
+import { g as A, a as k, b as C, c as N, e as R } from "./axis_util-Bu4h7XWV.js";
+import { t as U, m as W } from "./shared-Ca6iDobD.js";
+import { j as _, f as y } from "./gpgpu_math-CNslybmD.js";
+import { g as G, b as L } from "./kernel_funcs_utils-CUxJCg0g.js";
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+class w {
+  constructor(s, e) {
+    this.variableNames = ["x"];
+    const { windowSize: t, batchSize: n, inSize: l, outSize: r } = s;
+    this.outputShape = [n, r];
+    const i = Math.floor(t / 4) * 4, c = t % 4;
+    let o = "sumValue += dot(values, ones);";
+    if (e != null) {
+      const p = 1 / e;
+      o = `sumValue += dot(values * ${T(p) ? p.toPrecision(2) : p}, ones);`;
+    }
+    let u = "";
+    l % t > 0 && (u = `
+        if (inIdx < 0 || inIdx >= ${l}) {
+          return 0.0;
+        }
+      `), this.userCode = `
+      const vec4 ones = vec4(1.0, 1.0, 1.0, 1.0);
+
+      float getValue(int batch, int inIdx) {
+        ${u}
+        return getX(batch, inIdx);
+      }
+
+      void main() {
+        ivec2 coords = getOutputCoords();
+        int batch = coords[0];
+        int outIdx = coords[1];
+        int inOffset = outIdx * ${t};
+
+        float sumValue = 0.0;
+
+        for (int i = 0; i < ${i}; i += 4) {
+          int inIdx = inOffset + i;
+          vec4 values = vec4(
+            getValue(batch, inIdx),
+            getValue(batch, inIdx + 1),
+            getValue(batch, inIdx + 2),
+            getValue(batch, inIdx + 3)
+          );
+
+          ${o}
+        }
+
+        int inIdx = inOffset + ${i};
+        if (${c === 1}) {
+          vec4 values = vec4(getValue(batch, inIdx), 0.0, 0.0, 0.0);
+
+          ${o}
+        } else if (${c === 2}) {
+          vec4 values = vec4(
+            getValue(batch, inIdx),
+            getValue(batch, inIdx + 1), 0.0, 0.0);
+
+          ${o}
+        } else if (${c === 3}) {
+          vec4 values = vec4(
+            getValue(batch, inIdx),
+            getValue(batch, inIdx + 1),
+            getValue(batch, inIdx + 2), 0.0);
+
+          ${o}
+        }
+        setOutput(sumValue);
+      }
+    `;
+  }
+}
+/**
+ * @license
+ * Copyright 2017 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+class X {
+  constructor(s, e) {
+    this.variableNames = ["x"];
+    const { windowSize: t, batchSize: n, inSize: l, outSize: r } = s;
+    this.outputShape = [n, r];
+    let i = "0.0", c = "";
+    e === "prod" ? i = "1.0" : e === "min" ? (i = "1.0 / 1e-20", c = "min") : e === "max" && (i = "-1.0 / 1e-20", c = "max");
+    let o = `${e}(${e}(${e}(minMaxValue[0], minMaxValue[1]), minMaxValue[2]), minMaxValue[3])`;
+    e === "sum" ? o = "sumValue" : e === "prod" ? o = "prodValue" : e === "all" ? o = "allValue" : e === "any" && (o = "anyValue");
+    const u = Math.floor(t / 4) * 4, p = t % 4;
+    let h = `
+      if (${e === "sum"}) {
+        sumValue += dot(values, ones);
+      } else if (${e === "prod"}) {
+        vec2 tmp = vec2(values[0], values[1]) * vec2(values[2], values[3]);
+        prodValue *= tmp[0] * tmp[1];
+      } else {
+        minMaxValue = ${c}(values, minMaxValue);
+        if (${e === "min"} || ${e === "max"}) {
+          minMaxValue = ${c}(values, minMaxValue);
+          bvec4 isNaN = isnan(values);
+          if (isNaN.r || isNaN.g || isNaN.b || isNaN.a) {
+            minMaxValue = vec4(NAN);
+          }
+        }
+      }
+    `, d = "vec4";
+    e === "all" ? (i = "1.0", h = `
+        bool reducedAllValue = all(values);
+        float floatedReducedAllValue = float(reducedAllValue);
+        allValue = float(allValue >= 1.0 && floatedReducedAllValue >= 1.0);
+      `, d = "bvec4") : e === "any" && (i = "0.0", h = `
+        bool reducedAnyValue = any(values);
+        float floatedReducedAnyValue = float(reducedAnyValue);
+        anyValue = float(anyValue >= 1.0 || floatedReducedAnyValue >= 1.0);
+      `, d = "bvec4");
+    let f = "";
+    l % t > 0 && (f = `
+        if (inIdx < 0 || inIdx >= ${l}) {
+          return initializationValue;
+        }
+      `), this.userCode = `
+      const float initializationValue = ${i};
+      const vec4 ones = vec4(1.0, 1.0, 1.0, 1.0);
+
+      float getValue(int batch, int inIdx) {
+        ${f}
+        return getX(batch, inIdx);
+      }
+
+      void main() {
+        ivec2 coords = getOutputCoords();
+        int batch = coords[0];
+        int outIdx = coords[1];
+        int inOffset = outIdx * ${t};
+
+        vec4 minMaxValue = vec4(${i});
+        float prodValue = 1.0;
+        float sumValue = 0.0;
+        float allValue = 1.0;
+        float anyValue = 0.0;
+
+        for (int i = 0; i < ${u}; i += 4) {
+          int inIdx = inOffset + i;
+          ${d} values = ${d}(
+            getValue(batch, inIdx),
+            getValue(batch, inIdx + 1),
+            getValue(batch, inIdx + 2),
+            getValue(batch, inIdx + 3)
+          );
+
+          ${h}
+        }
+
+        int inIdx = inOffset + ${u};
+        if (${p === 1}) {
+          ${d} values = ${d}(
+            getValue(batch, inIdx),
+            initializationValue,
+            initializationValue,
+            initializationValue
+          );
+
+          ${h}
+        } else if (${p === 2}) {
+          ${d} values = ${d}(
+            getValue(batch, inIdx),
+            getValue(batch, inIdx + 1),
+            initializationValue,
+            initializationValue
+          );
+
+          ${h}
+        } else if (${p === 3}) {
+          ${d} values = ${d}(
+            getValue(batch, inIdx),
+            getValue(batch, inIdx + 1),
+            getValue(batch, inIdx + 2),
+            initializationValue
+          );
+
+          ${h}
+        }
+        setOutput(${o});
+      }
+    `;
+  }
+}
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function q(a) {
+  const s = [];
+  for (; s.length === 0 || s[s.length - 1].outSize !== 1; ) {
+    const e = s.length ? s[s.length - 1].outSize : a[1], t = _(e);
+    s.push({
+      inSize: e,
+      windowSize: t,
+      outSize: Math.ceil(e / t)
+    });
+  }
+  return s;
+}
+function M(a, s, e, t) {
+  const n = q(a.shape);
+  let l = a;
+  for (let r = 0; r < n.length; r++) {
+    const { inSize: i, windowSize: c, outSize: o } = n[r];
+    let u, p;
+    e === "mean" ? u = r === 0 ? new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, i) : new w({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }) : u = new X({ windowSize: c, inSize: i, batchSize: a.shape[0], outSize: o }, e), p = l, l = t.runWebGLProgram(u, [l], s), p.dataId !== a.dataId && t.disposeIntermediateTensorInfo(p);
+  }
+  return l;
+}
+/**
+ * @license
+ * Copyright 2017 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+class Y {
+  constructor(s, e) {
+    this.variableNames = ["A"];
+    const t = new Array(s.length);
+    for (let r = 0; r < t.length; r++)
+      t[r] = s[e[r]];
+    this.outputShape = t, this.rank = t.length;
+    const n = y(this.rank), l = H(e);
+    this.userCode = `
+    void main() {
+      ${n} resRC = getOutputCoords();
+      setOutput(getA(${l}));
+    }
+    `;
+  }
+}
+function H(a) {
+  const s = a.length;
+  if (s > 6)
+    throw Error(`Transpose for rank ${s} is not yet supported`);
+  const e = ["resRC.x", "resRC.y", "resRC.z", "resRC.w", "resRC.u", "resRC.v"], t = new Array(s);
+  for (let n = 0; n < a.length; n++)
+    t[a[n]] = e[n];
+  return t.join();
+}
+/**
+ * @license
+ * Copyright 2019 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+class J {
+  constructor(s, e) {
+    this.variableNames = ["A"], this.packedInputs = !0, this.packedOutput = !0;
+    const t = new Array(s.length);
+    for (let u = 0; u < t.length; u++)
+      t[u] = s[e[u]];
+    if (this.outputShape = t, this.rank = t.length, this.rank > 6)
+      throw Error(`Packed transpose for rank ${this.rank} is not yet supported.`);
+    const n = y(this.rank), l = G("rc", this.rank), r = new Array(this.rank);
+    for (let u = 0; u < e.length; u++)
+      r[e[u]] = l[u];
+    const i = `vec2(${r.slice(-2).join()})`, c = `++${l[this.rank - 1]} < ${t[this.rank - 1]}`, o = `getChannel(getA(${r.join()}), ${i})`;
+    this.userCode = `
+    void main() {
+      ${n} rc = getOutputCoords();
+      vec4 result = vec4(0.);
+      result[0] = ${o};
+      if(${c}) {
+        result[1] = ${o};
+      }
+      --${l[this.rank - 1]};
+      if(++${l[this.rank - 2]} < ${t[this.rank - 2]}) {
+        result[2] = ${o};
+        if(${c}) {
+          result[3] = ${o};
+        }
+      }
+      setOutput(result);
+    }
+    `;
+  }
+}
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function P(a, s, e) {
+  const t = E().getBool("WEBGL_PACK_ARRAY_OPERATIONS") ? new J(a.shape, s) : new Y(a.shape, s);
+  return e.runWebGLProgram(t, [a], a.dtype);
+}
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function Q(a, s, e, t) {
+  const n = s, l = a.shape.length, r = O(n, a.shape);
+  let i = r;
+  const c = A(i, l), o = c != null;
+  let u = a;
+  o && (u = P(a, c, t), i = k(i.length, l)), C("sum", i, l);
+  const [p, h] = N(u.shape, i);
+  let d = p;
+  e && (d = R(p, r));
+  const f = V(h), g = V(a.shape) / f, x = $({ inputs: { x: u }, attrs: { shape: [g, f] }, backend: t }), b = B(a.dtype), I = M(x, b, "sum", t), m = $({ inputs: { x: I }, attrs: { shape: d }, backend: t });
+  return t.disposeIntermediateTensorInfo(x), t.disposeIntermediateTensorInfo(I), o && t.disposeIntermediateTensorInfo(u), m;
+}
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function Z(a) {
+  const { inputs: s, backend: e, attrs: t } = a, { x: n } = s, { axis: l, keepDims: r } = t;
+  return Q(n, l, r, e);
+}
+const de = {
+  kernelName: F,
+  backendName: "webgl",
+  kernelFunc: Z
+};
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function ee(a, s, e, t) {
+  const n = V(s), r = V(a.shape) / n, i = $({ inputs: { x: a }, attrs: { shape: [r, n] }, backend: t }), c = M(i, a.dtype, "max", t), o = $({ inputs: { x: c }, attrs: { shape: e }, backend: t });
+  return t.disposeIntermediateTensorInfo(i), t.disposeIntermediateTensorInfo(c), o;
+}
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function te(a) {
+  const { inputs: s, backend: e, attrs: t } = a, { x: n } = s, { reductionIndices: l, keepDims: r } = t, i = n.shape.length, c = O(l, n.shape);
+  let o = c;
+  const u = A(o, i), p = u != null, h = e.shouldExecuteOnCPU([n]);
+  let d = n;
+  if (p) {
+    if (h) {
+      const I = e.texData.get(d.dataId).values, m = new Array(i);
+      for (let v = 0; v < m.length; v++)
+        m[v] = n.shape[u[v]];
+      const z = U(I, n.shape, n.dtype, u, m);
+      d = e.makeTensorInfo(m, n.dtype);
+      const D = e.texData.get(d.dataId);
+      D.values = z;
+    } else
+      d = P(n, u, e);
+    o = k(o.length, i);
+  }
+  C("max", o, i);
+  const [f, S] = N(d.shape, o);
+  let g = f;
+  r && (g = R(f, c));
+  let x;
+  if (h) {
+    const I = e.texData.get(d.dataId).values, m = W(I, V(S), g, n.dtype);
+    x = e.makeTensorInfo(g, n.dtype);
+    const z = e.texData.get(x.dataId);
+    z.values = m;
+  } else
+    x = ee(d, S, g, e);
+  return p && e.disposeIntermediateTensorInfo(d), x;
+}
+const pe = {
+  kernelName: j,
+  backendName: "webgl",
+  kernelFunc: te
+};
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+const ae = `
+if (a == b) {
+  return 1.0;
+};
+return a / b;`, se = `
+  // vec4 one = vec4(equal(a, b));
+  // return one + (vec4(1.0) - one) * a / b;
+  vec4 result = a / b;
+  if(a.x == b.x) {
+    result.x = 1.;
+  }
+  if(a.y == b.y) {
+    result.y = 1.;
+  }
+  if(a.z == b.z) {
+    result.z = 1.;
+  }
+  if(a.w == b.w) {
+    result.w = 1.;
+  }
+
+  return result;
+`, ne = L({ opSnippet: ae, packedOpSnippet: se, checkOutOfBounds: !0 }), he = {
+  kernelName: K,
+  backendName: "webgl",
+  kernelFunc: ne
+};
+export {
+  M as a,
+  pe as b,
+  he as c,
+  de as d,
+  te as m,
+  ne as r,
+  Z as s,
+  P as t
+};

package/dist/Reshape-t7Kcikjk.js

@@ -0,0 +1,127 @@
+import { g as c, aa as C, i as f, D as R } from "./index-BAzbokzv.js";
+import { u as g, g as I, a as x, b as F, c as $, d as u, e as m, i as l } from "./gpgpu_math-CNslybmD.js";
+/**
+ * @license
+ * Copyright 2018 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+class S {
+  constructor(t, i) {
+    this.variableNames = ["A"], this.packedInputs = !0, this.packedOutput = !0, this.customUniforms = [{ name: "inputShape", type: "ivec3" }], this.outputShape = t, this.enableShapeUniforms = g(this.outputShape.length);
+    let a = "";
+    for (let e = 0; e < 4; e++) {
+      let o = "thisRC = rc;";
+      e % 2 === 1 && (o += "thisRC.z += 1;"), e > 1 && (o += "thisRC.y += 1;"), a += `
+        ${o}
+        ${e > 0 ? "if(thisRC.y < rows && thisRC.z < cols){" : ""}
+          int flatIndex = getFlatIndex(thisRC);
+
+          ivec3 inputRC = inputCoordsFromReshapedOutCoords(flatIndex);
+          vec2 inputRCInnerDims = vec2(float(inputRC.y),float(inputRC.z));
+
+          result[${e}] =
+            getChannel(getA(inputRC.x, inputRC.y, inputRC.z), inputRCInnerDims);
+        ${e > 0 ? "}" : ""}
+      `;
+    }
+    this.userCode = `
+      ${v(i, this.enableShapeUniforms)}
+      ${this.enableShapeUniforms ? I() : x(t)}
+
+      void main() {
+        ivec3 rc = getOutputCoords();
+
+        vec4 result = vec4(0.);
+
+        ivec3 thisRC;
+        int rows = ${this.enableShapeUniforms ? "outShape[1]" : t[1]};
+        int cols = ${this.enableShapeUniforms ? "outShape[2]" : t[2]};
+
+        ${a}
+
+        setOutput(result);
+      }
+    `;
+  }
+}
+function v(s, t) {
+  return `
+    ivec3 inputCoordsFromReshapedOutCoords(int index) {
+      ${t ? F(["r", "c", "d"], "inputShape") : $(["r", "c", "d"], s)}
+      return ivec3(r, c, d);
+    }
+  `;
+}
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function b(s, t, i) {
+  const a = [
+    u(s.shape),
+    ...m(s.shape)
+  ], e = {
+    dtype: s.dtype,
+    shape: a,
+    dataId: s.dataId
+  }, o = [
+    u(t),
+    ...m(t)
+  ], r = new S(o, a), p = !0, n = [a], h = i.runWebGLProgram(r, [e], s.dtype, n, p);
+  return { dataId: h.dataId, shape: t, dtype: h.dtype };
+}
+/**
+ * @license
+ * Copyright 2020 Google LLC. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+function y(s) {
+  const { inputs: t, backend: i, attrs: a } = s, { x: e } = t, { shape: o } = a, r = i, p = c(e.shape), n = C(o, p), h = c(n);
+  f(p === h, () => `The new shape (${n}) has ${h} elements and the old shape (${e.shape}) has ${p} elements. The new shape and old shape must have the same number of elements.`);
+  const d = r.texData.get(e.dataId);
+  return d.isPacked && !l(e.shape, n) && !(d.texture !== null && l(d.shape, n)) ? b(e, n, r) : (r.incRef(e.dataId), { dataId: e.dataId, shape: n, dtype: e.dtype });
+}
+const U = {
+  kernelName: R,
+  backendName: "webgl",
+  kernelFunc: y
+};
+export {
+  S as R,
+  U as a,
+  y as r
+};

package/dist/TeachableLLM.d.ts

@@ -18,6 +18,7 @@ export default class TeachableLLM {
     private _model?;
     private _tokeniser?;
     private _status;
+    private _memoryRequirements?;
     meta: TeachableLLMMeta;
     constructor(tokeniser?: ITokeniser, model?: NanoGPT);
     get vocab(): string[];
@@ -27,6 +28,7 @@ export default class TeachableLLM {
     get tokeniser(): ITokeniser;
     get status(): TeachableLLMStatus;
     get ready(): boolean;
+    estimateTrainingMemoryUsage(batchSize: number): number;
     private setStatus;
     saveModel(options?: SaveOptions): Promise<Blob>;
     static loadModel(data: Blob | Buffer | string): TeachableLLM;