import { m as b, s as I, r as k } from "../../RealDiv-ioj6Z-ox.js";
import { r as v } from "../../Reshape-BZC-ebeR.js";
import { a as w, ah as P } from "../../index-DSGwv2Yx.js";
import { e as S } from "../../axis_util-QWWgLjut.js";
class T {
  variableNames = ["logits", "maxLogits"];
  outputShape;
  userCode;
  constructor(t) {
    this.outputShape = t, this.userCode = `
        void main() {
        ivec4 coords = getOutputCoords(); // [batch, nh, t1, t2]
            int b = coords.x;
            int h = coords.y;
            int t1 = coords.z;
            int t2 = coords.w;
            float x = getLogitsAtOutCoords();
            float maxLogit = getMaxLogits(b, h, t1);
            setOutput(exp(x - maxLogit));
        }
        `;
  }
}
class C {
  variableNames = ["exp", "sum"];
  outputShape;
  userCode;
  customUniforms = [
    { name: "dropoutRate", type: "float" },
    { name: "seed", type: "float" }
  ];
  constructor(t) {
    this.outputShape = t, this.userCode = `
        float random(ivec4 coords) {
            float x = float(coords.x * 4096 + coords.y * 256 + coords.z * 16 + coords.w);
            return fract(sin(seed + x) * 43758.5453123);
        }
        void main() {
            ivec4 coords = getOutputCoords();
            float numerator = getExp(coords.x, coords.y, coords.z, coords.w);
            float denominator = getSum(coords.x, coords.y, coords.z, coords.w);
            float val = numerator / denominator;
            float keepProb = 1.0 - dropoutRate;
            float rand = random(coords);
            float mask = step(rand, keepProb);
            setOutput(val * mask / keepProb);
        }
        `;
  }
}
function L(r) {
  const { inputs: t, attrs: m } = r, { logits: e } = t, { dim: u, dropoutRate: n, seed: c } = m, o = r.backend;
  if (!e)
    throw new Error("Error in softmax: input logits is null");
  const i = P([u], e.shape), d = b({
    inputs: { x: e },
    backend: o,
    attrs: { reductionIndices: i, keepDims: !1 }
  }), f = S(d.shape, i), l = new T(e.shape), s = o.runWebGLProgram(l, [e, d], "float32");
  o.disposeIntermediateTensorInfo(d);
  const p = I({ inputs: { x: s }, backend: o, attrs: { axis: i, keepDims: !1 } }), a = v({ inputs: { x: p }, backend: o, attrs: { shape: f } });
  if (n !== void 0 && n > 0) {
    const h = new C(e.shape), g = o.runWebGLProgram(h, [s, a], "float32", [
      [n],
      [c ?? Math.random() * 1e4]
    ]);
    return o.disposeIntermediateTensorInfo(s), o.disposeIntermediateTensorInfo(p), o.disposeIntermediateTensorInfo(a), g;
  }
  const x = k({ inputs: { a: s, b: a }, backend: o });
  return o.disposeIntermediateTensorInfo(s), o.disposeIntermediateTensorInfo(p), o.disposeIntermediateTensorInfo(a), x;
}
const E = {
  kernelName: "FusedSoftmax",
  backendName: "webgl",
  kernelFunc: L
};
w(E);
export {
  L as softmax
};
