import { e as u } from "../../webgpu_program-CbjdYLYk.js";
import { f as h, c as p } from "../../webgpu_util-DuofJBMo.js";
import { a as c, a6 as r } from "../../index-DSGwv2Yx.js";
class l {
  variableNames = ["labels", "logits", "values"];
  outputShape;
  shaderKey = "GatherSub";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  size = !0;
  constructor(t) {
    this.outputShape = [t], this.dispatchLayout = h(this.outputShape), this.dispatch = p(this.dispatchLayout, this.outputShape, this.workgroupSize);
  }
  getUserCode() {
    return `
        ${u("index")} {
            if (index < uniforms.size) {
                let idx = i32(labels[index]);
                let val = values[index];
                let logit = logits[index * uniforms.logitsShape[1] + idx];
                setOutputAtIndex(index, val - logit);
            }
        }
    `;
  }
}
function d(e) {
  const { logits: t, labels: a, values: s } = e.inputs, o = e.backend, i = a.shape[0];
  r(s.shape, [i], "Error in EfficientGatherSub: "), r(a.shape, [i], "Error in EfficientGatherSub: ");
  const n = new l(i);
  return o.runWebGPUProgram(n, [a, t, s], "float32");
}
const f = {
  kernelName: "EfficientGatherSub",
  backendName: "webgpu",
  kernelFunc: d
};
c(f);
