@genai-fi/nanogpt 0.9.0 → 0.10.0

package/dist/ops/webgpu/slice16.js

@@ -0,0 +1,71 @@
+import { b as c, c as f, e as l } from "../../webgpu_program-Cigz-7RF.js";
+import { f as g, c as S } from "../../webgpu_util-BBCnKm2X.js";
+import "../../index-ZyQhjEPo.js";
+import { r as k } from "../../tensor_util-DV-FP5Q3.js";
+import { p as y, a as $ } from "../../slice_util-DtEldBfK.js";
+import { s as x } from "../../tensor-DdQUJZlz.js";
+function b(o) {
+  switch (o) {
+    case 1:
+      return "1D";
+    case 2:
+      return "2D";
+    case 3:
+      return "3D";
+    case 4:
+      return "4D";
+  }
+  return "";
+}
+class w {
+  variableNames = ["source"];
+  uniforms;
+  outputShape;
+  shaderKey;
+  rank;
+  dispatchLayout;
+  dispatch;
+  workPerThread = 1;
+  workgroupSize = [64, 1, 1];
+  start;
+  size = !0;
+  constructor(e, t) {
+    this.outputShape = t, this.rank = t.length, this.dispatchLayout = g(this.outputShape), this.dispatch = S(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+      this.workPerThread,
+      1,
+      1
+    ]), this.start = e, this.uniforms = `start : ${c(e.length)}, `, this.shaderKey = "slice";
+  }
+  getUserCode() {
+    const e = c(this.rank);
+    let t;
+    return this.start.length === 1 ? t = this.outputShape.map(() => "sourceLoc = uniforms.start + coords;") : t = this.outputShape.map((r, s) => `sourceLoc.${p[s]} = uniforms.start.${f(s)} + coords.${p[s]};`), `
+      ${l("index")} {
+        if (index < uniforms.size) {
+          var sourceLoc : ${e};
+          let coords = getCoordsFromIndex(index);
+          ${t.join(`
+`)}
+          result[index] = source[getIndexFromCoords${b(this.rank)}(sourceLoc, uniforms.sourceShape)];
+        }
+      }
+    `;
+  }
+}
+const p = ["x", "y", "z", "w", "u", "v"];
+function C(o) {
+  const { inputs: e, backend: t, attrs: n } = o, { x: r } = e, { begin: s, size: h } = n, [i, a] = y(r, s, h);
+  if ($(r, i, a), x(a) === 0)
+    return t.makeTensorInfo(a, r.dtype, []);
+  const d = new w(i, a), m = [{ type: "int32", data: i }], u = t.runWebGPUProgram(d, [r], r.dtype, m);
+  return u.packed = !0, u;
+}
+const D = {
+  kernelName: "Slice16",
+  backendName: "webgpu",
+  kernelFunc: C
+};
+k(D);
+export {
+  C as slice
+};

package/dist/{variable-Bm2OFwGI.js → ops/webgpu/softmax16.d.ts}

@@ -1,7 +1,6 @@
-import { G as i } from "./index-BzFyqcy-.js";
 /**
  * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
+ * Copyright 2023 Google LLC.
  * 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
@@ -15,9 +14,4 @@ import { G as i } from "./index-BzFyqcy-.js";
  * limitations under the License.
  * =============================================================================
  */
-function m(r, a = !0, e, t) {
-  return i.makeVariable(r, a, e, t);
-}
-export {
-  m as v
-};
+export {};

package/dist/ops/webgpu/softmax16.js

@@ -0,0 +1,23 @@
+import { e as S } from "../../index-ZyQhjEPo.js";
+import { reshape16 as h } from "../reshape16.js";
+import b from "./softmax16_program.js";
+import d from "./softmax16_subgroup_program.js";
+import x from "./utils/deviceInfo.js";
+import { r as k } from "../../tensor_util-DV-FP5Q3.js";
+import { r as l } from "../../reshape-DevtBWtf.js";
+import { s as z } from "../../tensor-DdQUJZlz.js";
+function I(a) {
+  const { inputs: m, backend: e, attrs: p } = a, { logits: o } = m, { dim: s } = p, i = e.subgroupMinSize, c = e.subgroupMaxSize, u = x(e).subgroupsSupported, r = l(o, [
+    z(o.shape) / o.shape[s],
+    o.shape[s]
+  ]), f = u ? new d(r.shape, i, c) : new b(r.shape), n = e.runWebGPUProgram(f, [r], "int32");
+  n.packed = !0, r.dispose();
+  const t = S().makeTensorFromTensorInfo(n), g = h(t, o.shape);
+  return t.dispose(), g;
+}
+const P = {
+  kernelName: "Softmax16",
+  backendName: "webgpu",
+  kernelFunc: I
+};
+k(P);

package/dist/ops/webgpu/softmax16_program.d.ts

@@ -0,0 +1,13 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class SoftmaxProgram implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    constructor(outputShape: number[]);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/softmax16_program.js

@@ -0,0 +1,73 @@
+import { e } from "../../webgpu_program-Cigz-7RF.js";
+import { f as o } from "../../webgpu_util-BBCnKm2X.js";
+class i {
+  variableNames = ["logits"];
+  outputShape;
+  shaderKey;
+  dispatchLayout;
+  dispatch;
+  workgroupSize;
+  constructor(r) {
+    this.outputShape = r, this.dispatchLayout = o(this.outputShape), this.dispatch = [this.outputShape[0], 1, 1], this.outputShape[1] >= 4096 ? this.workgroupSize = [256, 1, 1] : this.outputShape[1] < 64 ? this.workgroupSize = [32, 1, 1] : this.workgroupSize = [64, 1, 1], this.shaderKey = "softmax16";
+  }
+  getUserCode() {
+    return `
+        var<workgroup> buf : array<f32, ${this.workgroupSize[0]}>;
+        const blockSize = ${this.workgroupSize[0]};
+        ${e("index")} {
+            let row = index / blockSize;
+            let tid = i32(localId.x);
+            let cols = uniforms.outShape[1];
+            let rowIdx = row * cols;
+
+            var threadMax = -3.402823e+38f;
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                threadMax = max(threadMax, max(value.x, value.y));
+            }
+            buf[tid] = threadMax;
+            workgroupBarrier();
+
+            for (var currSize = blockSize >> 1;  currSize > 0; currSize = currSize >> 1) {
+                if (tid < currSize) {
+                    buf[tid] = max(buf[tid], buf[tid + currSize]);
+                }
+                workgroupBarrier();
+            }
+
+            let rowMaxShared: f32 = buf[0];
+            workgroupBarrier();
+
+            var threadSum = 0.0f;
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                let subExp = exp(value.x - rowMaxShared);
+                threadSum += subExp;
+                let subExpY = exp(value.y - rowMaxShared);
+                threadSum += subExpY;
+            }
+            buf[tid] = threadSum;
+            workgroupBarrier();
+
+            for (var currSize = blockSize >> 1;  currSize > 0; currSize = currSize >> 1) {
+                if (tid < currSize) {
+                    buf[tid] = buf[tid] + buf[tid + currSize];
+                }
+                workgroupBarrier();
+            }
+
+            let rowSumShared: f32 = buf[0];
+
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                let value1: f32 = exp(value.x - rowMaxShared) / rowSumShared;
+                let value2: f32 = exp(value.y - rowMaxShared) / rowSumShared;
+                result[rowIdx + col] = i32(pack2x16float(vec2<f32>(value1, value2)));
+            }
+        }
+      `;
+  }
+}
+export {
+  i as default
+};

package/dist/ops/webgpu/softmax16_subgroup_program.d.ts

@@ -0,0 +1,17 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class SoftmaxSubgroupProgram implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    minSubgroupSize: number;
+    maxSubgroupSize: number;
+    subgroups: boolean;
+    subgroupBuiltins: boolean;
+    constructor(outputShape: number[], minSubgroupSize: number, maxSubgroupSize: number);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/softmax16_subgroup_program.js

@@ -0,0 +1,75 @@
+import { e as o } from "../../webgpu_program-Cigz-7RF.js";
+import { f as u } from "../../webgpu_util-BBCnKm2X.js";
+class i {
+  variableNames = ["logits"];
+  outputShape;
+  shaderKey;
+  dispatchLayout;
+  dispatch;
+  workgroupSize;
+  minSubgroupSize;
+  maxSubgroupSize;
+  subgroups = !0;
+  subgroupBuiltins = !1;
+  constructor(e, a, r) {
+    this.minSubgroupSize = a, this.maxSubgroupSize = r, this.outputShape = e, this.dispatchLayout = u(this.outputShape), this.dispatch = [this.outputShape[0], 1, 1], this.workgroupSize = [Math.min(64, r), 1, 1], a !== r && (this.subgroupBuiltins = !0, this.workgroupSize = [64, 1, 1]), this.shaderKey = "softmax16subgroup";
+  }
+  getUserCode() {
+    const e = this.maxSubgroupSize !== this.minSubgroupSize;
+    return `
+        ${e ? `var<workgroup> bestValues : array<f32, ${this.workgroupSize[0]}>;` : ""}
+        const blockSize = ${this.workgroupSize[0]};
+        ${o("index")} {
+            let row = index / blockSize;
+            let tid = i32(localId.x);
+            let cols = uniforms.outShape[1];
+            let rowIdx = row * cols;
+
+            var threadMax = -3.402823e+38f;
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                threadMax = max(threadMax, max(value.x, value.y));
+            }
+
+            threadMax = subgroupMax(threadMax);
+            ${e ? `
+                let lane = localId.x % subgroupSize;
+                if (lane == 0) {
+                    bestValues[localId.x / subgroupSize] = threadMax;
+                }
+                workgroupBarrier();
+                let numSubgroups = blockSize / subgroupSize;
+                threadMax = select(-3.402823e+38f, bestValues[lane], lane < numSubgroups);
+                threadMax = subgroupMax(threadMax);
+                workgroupBarrier();    
+            ` : ""}
+
+            var threadSum = 0.0f;
+            for (var col = tid; col < cols; col += blockSize) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                let subExp = exp(value - threadMax);
+                threadSum += subExp.x + subExp.y;
+            }
+
+            threadSum = subgroupAdd(threadSum);
+            ${e ? `
+                if (lane == 0) {
+                    bestValues[localId.x / subgroupSize] = threadSum;
+                }
+                workgroupBarrier();
+                threadSum = select(0.0f, bestValues[lane], lane < numSubgroups);
+                threadSum = subgroupAdd(threadSum);    
+            ` : ""}
+
+            for (var col = tid; col < cols; col += ${e ? "i32(subgroupSize)" : "blockSize"}) {
+                let value = unpack2x16float(u32(logits[rowIdx + col]));
+                let valuePair: vec2<f32> = exp(value - threadMax) / threadSum;
+                result[rowIdx + col] = i32(pack2x16float(valuePair));
+            }
+        }
+      `;
+  }
+}
+export {
+  i as default
+};

package/dist/ops/webgpu/softmax16grad.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/softmax16grad.js

@@ -0,0 +1,38 @@
+import "../../index-ZyQhjEPo.js";
+import { createReduceInfo as p, reduce as m, ReduceProgram as i } from "./utils/reductions.js";
+import { isPackedTensor as n } from "../../utilities/packed.js";
+import l from "./utils/deviceInfo.js";
+import { r as k } from "../../tensor_util-DV-FP5Q3.js";
+class x extends i {
+  constructor(e, t) {
+    super(e, t, { reductionOp: "sum", elementwise: !0 }, !0), this.shaderKey = "SoftmaxGrad16", this.variableNames = ["dy", "softmaxOutput"], this.variableComponents = [1, 1];
+  }
+  getReadSnippet() {
+    return `
+            let d: vec2<f32> = unpack2x16float(u32(dy[index]));
+            let l: vec2<f32> = unpack2x16float(u32(softmaxOutput[index]));
+            return d * l;
+        `;
+  }
+  getWriteSnippet() {
+    return `
+            let d: vec2<f32> = unpack2x16float(u32(dy[offset + k]));
+            let l: vec2<f32> = unpack2x16float(u32(softmaxOutput[offset + k]));
+            let outVal = l * (d - bestValue);
+            result[offset + k] = i32(pack2x16float(outVal));
+        `;
+  }
+}
+function b(o) {
+  const { dy: e, softmaxOutput: t } = o.inputs, r = o.backend, s = l(r), u = n(e), c = n(t);
+  if (!(u && c))
+    throw new Error("softmaxGradGPU: dy and softmaxOutput must be packed tensors");
+  const a = [e, t], f = p(a, -1), d = new x(s, f);
+  return m(d, a, r);
+}
+const v = {
+  kernelName: "Softmax16Grad",
+  backendName: "webgpu",
+  kernelFunc: b
+};
+k(v);

package/dist/ops/webgpu/sub16.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/sub16.js

@@ -0,0 +1,14 @@
+import "../../index-ZyQhjEPo.js";
+import { BinaryOpProgram as p } from "./utils/binary_op.js";
+import { B as m } from "../../binary_op_util-pKXltfxI.js";
+import { r as s } from "../../tensor_util-DV-FP5Q3.js";
+function c(r) {
+  const { a: e, b: n } = r.inputs, t = r.backend, a = new p(m.SUB, e.shape, n.shape), o = t.runWebGPUProgram(a, [e, n], "int32");
+  return o.packed = !0, o;
+}
+const i = {
+  kernelName: "Sub16",
+  backendName: "webgpu",
+  kernelFunc: c
+};
+s(i);

package/dist/ops/webgpu/sum16.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/sum16.js

@@ -0,0 +1,40 @@
+import { createReduceInfo as g, reduce as h, ReduceProgram as x } from "./utils/reductions.js";
+import "../../index-ZyQhjEPo.js";
+import { isPackedTensor as k } from "../../utilities/packed.js";
+import { transpose16 as A } from "../transpose16.js";
+import P from "./utils/deviceInfo.js";
+import { r as b } from "../../tensor_util-DV-FP5Q3.js";
+import { s as I } from "../../sum-_fzj5ZTB.js";
+import { p as w } from "../../tensor-DdQUJZlz.js";
+import { a as D, b as K } from "../../axis_util-BvHEw88j.js";
+class v extends x {
+  shaderKey = "sum16";
+  constructor(e, o, t) {
+    super(
+      e,
+      o,
+      {
+        reductionOp: "sum",
+        elementwise: !1
+      },
+      t
+    ), t && (this.shaderKey += "_packed");
+  }
+}
+function y(r) {
+  const { x: e } = r.inputs, { axis: o, keepDims: t } = r.attrs, m = r.backend, a = [], p = P(m), c = k(e);
+  if (!c)
+    return I(e, o, t);
+  let n = w(o ?? -1, e.shape);
+  const i = D(n, e.shape.length);
+  let s = e;
+  i != null && (s = A(e, i), n = K(n.length, s.shape.length), a.push(s));
+  const u = g([s], -1), f = new v(p, u, c), d = h(f, [s], m);
+  return a.forEach((l) => l.dispose()), d;
+}
+const N = {
+  kernelName: "Sum16",
+  backendName: "webgpu",
+  kernelFunc: y
+};
+b(N);

package/dist/ops/webgpu/transpose16.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/transpose16.js

@@ -0,0 +1,35 @@
+import { isPackedTensor as d } from "../../utilities/packed.js";
+import { e as k } from "../../index-ZyQhjEPo.js";
+import { reshape16 as i } from "../reshape16.js";
+import l from "./transpose16_shared_program.js";
+import P from "./transpose16_program.js";
+import { r as b } from "../../tensor_util-DV-FP5Q3.js";
+import { t as T } from "../../transpose-DKELTqhe.js";
+function w(a) {
+  const { inputs: u, attrs: h } = a, { x: e } = u, { perm: r } = h, m = a.backend, c = d(e);
+  if (c && r[r.length - 1] !== e.shape.length - 1) {
+    const n = e.shape.length, t = n === 4 ? r.map((s) => s - 1).filter((s) => s >= 0) : r, p = n === 4 ? i(e, [e.shape[0] * e.shape[1], e.shape[2], e.shape[3]]) : e, f = new l(p.shape, t), o = m.runWebGPUProgram(f, [p], "int32");
+    if (o.packed = !0, n === 4) {
+      p.dispose();
+      const s = k().makeTensorFromTensorInfo(o), g = i(s, [
+        e.shape[0],
+        e.shape[1],
+        o.shape[1],
+        o.shape[2]
+      ]);
+      return s.dispose(), g;
+    }
+    return o;
+  }
+  if (c) {
+    const n = new P(e.shape, r), t = m.runWebGPUProgram(n, [e], "int32");
+    return t.packed = !0, t;
+  } else
+    return T(e, r);
+}
+const F = {
+  kernelName: "Transpose16",
+  backendName: "webgpu",
+  kernelFunc: w
+};
+b(F);

package/dist/ops/webgpu/transpose16_program.d.ts

@@ -0,0 +1,16 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+export default class TransposeProgram16 implements WebGPUProgram {
+    variableNames: string[];
+    shaderKey: string;
+    outputShape: number[];
+    dispatchLayout: {
+        x: number[];
+    };
+    dispatch: [number, number, number];
+    workPerThread: number;
+    workgroupSize: [number, number, number];
+    newDim: number[];
+    size: boolean;
+    constructor(aShape: number[], newDim: number[]);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/transpose16_program.js

@@ -0,0 +1,50 @@
+import { f as a, c as i } from "../../webgpu_util-BBCnKm2X.js";
+import { b as h, e as d, c as n } from "../../webgpu_program-Cigz-7RF.js";
+function p(r) {
+  const e = r.length;
+  if (e > 6)
+    throw Error(`Transpose for rank ${e} is not yet supported`);
+  const o = new Array(e);
+  for (let t = 0; t < r.length; t++)
+    o[r[t]] = `coords.${n(t)}`;
+  return o.join();
+}
+class l {
+  variableNames = ["A"];
+  shaderKey;
+  outputShape;
+  dispatchLayout;
+  dispatch;
+  workPerThread = 1;
+  workgroupSize = [64, 1, 1];
+  newDim;
+  size = !0;
+  constructor(e, o) {
+    const t = new Array(e.length);
+    for (let s = 0; s < t.length; s++)
+      t[s] = e[o[s]];
+    this.outputShape = t, this.dispatchLayout = a(this.outputShape), this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize, [
+      this.workPerThread,
+      1,
+      1
+    ]), this.newDim = o, this.shaderKey = `transpose16_${o}`;
+  }
+  getUserCode() {
+    const e = h(this.outputShape.length), o = p(this.newDim);
+    return `
+      ${d("index")} {
+        for(var i = 0; i < ${this.workPerThread}; i = i + 1) {
+          let flatIndex = index * ${this.workPerThread} + i;
+          if(flatIndex < uniforms.size) {
+            let coords = getCoordsFromIndex(flatIndex);
+            result[flatIndex] = A[getIndexFromCoords${this.outputShape.length}D(
+              ${e}(${o}), uniforms.aShape)];
+          }
+        }
+      }
+    `;
+  }
+}
+export {
+  l as default
+};

package/dist/ops/webgpu/transpose16_shared_program.d.ts

@@ -0,0 +1,15 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu/dist/webgpu_program';
+export default class TransposeSharedProgram16 implements WebGPUProgram {
+    variableNames: string[];
+    outputShape: number[];
+    shaderKey: string;
+    dispatchLayout: {
+        x: number[];
+        y: number[];
+        z?: number[];
+    };
+    dispatch: [number, number, number];
+    workgroupSize: [number, number, number];
+    constructor(aShape: number[], newDim: number[]);
+    getUserCode(): string;
+}

package/dist/ops/webgpu/transpose16_shared_program.js

@@ -0,0 +1,71 @@
+import { c as a } from "../../webgpu_util-BBCnKm2X.js";
+import { e as p } from "../../webgpu_program-Cigz-7RF.js";
+import "../../index-ZyQhjEPo.js";
+import { a as l } from "../../tensor-DdQUJZlz.js";
+class y {
+  variableNames = ["A"];
+  outputShape;
+  shaderKey;
+  dispatchLayout;
+  dispatch;
+  // Note that the maximum number of workgroup invocations by webgpu is 256.
+  // Nick: Reduce to 8x8
+  workgroupSize = [8, 8, 1];
+  constructor(t, o) {
+    const i = t.length, e = new Array(i), u = t.slice();
+    u[u.length - 1] *= 2;
+    for (let r = 0; r < e.length; r++)
+      e[r] = u[o[r]];
+    e[e.length - 1] /= 2, this.outputShape = e, this.dispatchLayout = i === 2 ? { x: [0], y: [1] } : { x: [1], y: [2], z: [0] }, this.dispatch = a(this.dispatchLayout, this.outputShape, this.workgroupSize, [2, 1, 1]), this.shaderKey = `transposeShared16_${i}`;
+  }
+  getUserCode() {
+    const t = this.outputShape.length;
+    l(
+      this.workgroupSize[0] === this.workgroupSize[1],
+      () => `Must be a square tile, current tile shape is ${this.workgroupSize[0]} x ${this.workgroupSize[1]}`
+    );
+    const o = this.workgroupSize[0] * 2;
+    return `
+      var<workgroup> tile : array<array<f32, ${o + 1}>, ${o}>;
+      ${p()} {
+        var x = i32(workgroupId.x) * ${o / 2} + i32(localId.x);
+        var y = i32(workgroupId.y) * ${o} + i32(localId.y);
+        let batch = ${t === 3 ? "i32(workgroupId.z)" : "0"};
+        let batchOffsetA = ${t === 3 ? "batch * uniforms.aShapeStrides[0]" : "0"};
+        let batchOffsetOut = ${t === 3 ? "batch * uniforms.outShapeStrides[0]" : "0"};
+
+        let inputWidth = uniforms.outShape[${t === 3 ? "1" : "0"}] / 2; // Output height
+        let inputHeight = uniforms.outShape[${t === 3 ? "2" : "1"}] * 2; // Output width
+        if (x < inputWidth && y < inputHeight) {
+            let unpackedA = unpack2x16float(u32(A[batchOffsetA + y * inputWidth + x]));
+            tile[localId.y][localId.x * 2] = unpackedA.x;
+            tile[localId.y][localId.x * 2 + 1] = unpackedA.y;
+        }
+        // Second load to cover the tile
+        y = y + ${this.workgroupSize[0]};
+        if (x < inputWidth && y < inputHeight) {
+            let unpackedA = unpack2x16float(u32(A[batchOffsetA + y * inputWidth + x]));
+            tile[localId.y + ${this.workgroupSize[0]}][localId.x * 2] = unpackedA.x;
+            tile[localId.y + ${this.workgroupSize[0]}][localId.x * 2 + 1] = unpackedA.y;
+        }
+        workgroupBarrier();
+
+        let outputWidth = uniforms.outShape[${t === 3 ? "2" : "1"}]; // Output width
+        let outputHeight = uniforms.outShape[${t === 3 ? "1" : "0"}] * 2; // Output height
+        x = i32(workgroupId.y) * ${o / 2} + i32(localId.x);
+        y = i32(workgroupId.x) * ${o} + i32(localId.y);
+        if (x < outputWidth && y < outputHeight) {
+          result[batchOffsetOut + y * outputWidth + x] = i32(pack2x16float(vec2<f32>(tile[localId.x * 2][localId.y], tile[localId.x * 2 + 1][localId.y])));
+        }
+        // Second store to cover the tile
+        y = y + ${this.workgroupSize[0]};
+        if (x < outputWidth && y < outputHeight) {
+          result[batchOffsetOut + y * outputWidth + x] = i32(pack2x16float(vec2<f32>(tile[localId.x * 2][localId.y + ${this.workgroupSize[0]}], tile[localId.x * 2 + 1][localId.y + ${this.workgroupSize[0]}])));
+        }
+      }
+    `;
+  }
+}
+export {
+  y as default
+};

package/dist/ops/webgpu/unpack16.d.ts

@@ -0,0 +1 @@
+export {};

package/dist/ops/webgpu/unpack16.js

@@ -0,0 +1,49 @@
+import { f as c, c as r } from "../../webgpu_util-BBCnKm2X.js";
+import { e as u } from "../../webgpu_program-Cigz-7RF.js";
+import "../../index-ZyQhjEPo.js";
+import { r as p } from "../../tensor_util-DV-FP5Q3.js";
+class l {
+  outputShape;
+  shaderKey = "Unpack16";
+  dispatchLayout;
+  dispatch;
+  workgroupSize = [64, 1, 1];
+  variableNames = ["x"];
+  size = !0;
+  uniforms;
+  outputComponent = 4;
+  variableComponents = [2];
+  scaling = !1;
+  constructor(t) {
+    this.outputShape = [...t.slice(0, -1), t[t.length - 1] * 2], this.dispatchLayout = c(this.outputShape), this.dispatch = r(this.dispatchLayout, this.outputShape, this.workgroupSize, [4, 1, 1]);
+  }
+  useScaling() {
+    this.shaderKey += "_Scaled", this.uniforms = "scaling : f32,", this.scaling = !0;
+  }
+  getUserCode() {
+    return `
+                ${u("index")} {
+                    let outIndex = index;
+                    if (outIndex < uniforms.size) {
+                        let xvec2 = x[index];
+                        let v1 = vec4<f32>(
+                            unpack2x16float(u32(xvec2.x)),
+                            unpack2x16float(u32(xvec2.y))
+                        ) ${this.scaling ? "* uniforms.scaling" : ""};
+                        result[outIndex] = v1;
+                    }
+                }`;
+  }
+}
+function h(e) {
+  const { x: t } = e.inputs, { scaling: a } = e.attrs, i = e.backend, n = new l(t.shape), s = a !== 1;
+  s && n.useScaling();
+  const o = [{ type: "float32", data: [1 / a] }];
+  return i.runWebGPUProgram(n, [t], "float32", s ? o : void 0);
+}
+const d = {
+  kernelName: "Unpack16",
+  backendName: "webgpu",
+  kernelFunc: h
+};
+p(d);

package/dist/ops/webgpu/utils/binary_op.d.ts

@@ -0,0 +1,19 @@
+import { WebGPUProgram } from '@tensorflow/tfjs-backend-webgpu';
+import { BinaryOpType } from '@tensorflow/tfjs-backend-webgpu/dist/binary_op_util';
+export { BinaryOpType };
+export declare class BinaryOpProgram implements WebGPUProgram {
+    dispatch: [number, number, number];
+    dispatchLayout: {
+        x: number[];
+    };
+    outputComponent: number;
+    op: BinaryOpType;
+    outputShape: number[];
+    shaderKey: string;
+    size: boolean;
+    variableNames: string[];
+    workgroupSize: [number, number, number];
+    variableComponents: number[];
+    constructor(op: BinaryOpType, aShape: number[], bShape: number[]);
+    getUserCode(): string;
+}