import { reduce as d, ReduceProgram as g } from "./utils/reductions.js";
import { a as S } from "../../index-DSGwv2Yx.js";
import h from "./utils/deviceInfo.js";
class k extends g {
  shaderKey = "clipscale";
  constructor(e, i, n) {
    super(
      e,
      i,
      {
        reductionOp: "sum",
        elementwise: !1,
        forceWorkgroupSize: n
      },
      !1
    ), this.uniforms += "scaling: f32, clipNorm: f32";
  }
  getReadSnippet() {
    return `
            return bitcast<f32>(u32(x[index]));
        `;
  }
  getWriteSnippet() {
    return `
            if (tid == 0) {
                let cnorm = uniforms.clipNorm;
                let gradNorm = sqrt(max(bestValue, 0.0));
                result[0] = (cnorm / max(cnorm, gradNorm)) * uniforms.scaling;
                result[1] = gradNorm;
            }
        `;
  }
}
function w(o) {
  const { x: e } = o.inputs, { invLossScaling: i, clipNorm: n } = o.attrs, s = o.backend, a = [];
  if (e.shape.length !== 1)
    throw new Error(`ClipScaleProgram requires 1D input, but got shape ${e.shape}`);
  const l = h(s);
  let r = 128, c = 1;
  const t = e.shape[0];
  t <= 16 ? r = 16 : t <= 32 ? r = 32 : t <= 64 && (r = 64), t > 128 && (c = Math.ceil(t / 128));
  const p = {
    inSize: r * c,
    outSize: 2,
    batchSize: 1,
    windowSize: r
  }, u = new k(l, p, r), m = d(u, [e], s, [
    { type: "float32", data: [i] },
    { type: "float32", data: [n] }
  ]);
  return a.forEach((f) => f.dispose()), m;
}
const b = {
  kernelName: "ClipScale",
  backendName: "webgpu",
  kernelFunc: w
};
S(b);
