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