import { a as C, t as R, e as I, N as G, _ as L, v as F, am as U } from "./index-DSGwv2Yx.js";
import { r as M } from "./Reshape-BZC-ebeR.js";
import { u as H } from "./gpgpu_math-3bCb5ooU.js";
import { m as B } from "./mat_mul-DP86qZtZ.js";
class W {
  constructor(e, s, a, n = !1, o = !1, r = !1, i = null, u = !1, l = !1) {
    this.variableNames = ["matrixA", "matrixB"], this.packedInputs = !0, this.packedOutput = !0, this.outputShape = a, this.enableShapeUniforms = H(this.outputShape.length);
    const p = n ? e[1] : e[2], h = Math.ceil(p / 2), d = n ? "i * 2, rc.y" : "rc.y, i * 2", x = o ? "rc.z, i * 2" : "i * 2, rc.z", b = n ? ["a.xxyy", "a.zzww"] : ["a.xxzz", "a.yyww"], m = o ? ["b.xzxz", "b.ywyw"] : ["b.xyxy", "b.zwzw"];
    let c = "", g = "";
    i && (u ? c = `vec4 activation(vec4 a) {
          vec4 b = getPreluActivationWeightsAtOutCoords();
          ${i}
        }` : l ? c = `vec4 activation(vec4 a) {
          vec4 b = getLeakyreluAlphaAtOutCoords();
          ${i}
        }` : c = `vec4 activation(vec4 x) {
          ${i}
        }`, g = "result = activation(result);");
    const $ = r ? "result += getBiasAtOutCoords();" : "";
    r && this.variableNames.push("bias"), u && this.variableNames.push("preluActivationWeights"), l && this.variableNames.push("leakyreluAlpha");
    let f = "rc.x", v = "rc.x";
    e[0] < s[0] ? f = `imod(rc.x, ${e[0]})` : s[0] < e[0] && (v = `imod(rc.x, ${s[0]})`), this.userCode = `
      ${c}
      // Don't use uniform for sharedDimensionPacked for performance.
      const float sharedDimension = ${h}.0;

      vec4 dot2x2ARowBCol(ivec3 rc) {
        vec4 result = vec4(0);
        int batchA = ${f};
        int batchB = ${v};
        for (int i = 0; i < ${h}; i++) {
          vec4 a = getMatrixA(batchA, ${d});
          vec4 b = getMatrixB(batchB, ${x});

          // These swizzled products need to be separately added.
          // See: https://github.com/tensorflow/tfjs/issues/1735
          result += (${b[0]} * ${m[0]});
          result += (${b[1]} * ${m[1]});
        }
        return result;
      }

      void main() {
        ivec3 rc = getOutputCoords();
        vec4 result = dot2x2ARowBCol(rc);

        ${$}

        ${g}

        setOutput(result);
      }
    `;
  }
}
const S = 0.7978845608028654, w = 0.044715, j = `
    vec4 x3 = x * x * x;
    vec4 inner = x + ${w} * x3;
    inner = ${S} * inner;
    inner = vec4(
        abs(inner[0]) > 15.0 ? sign(inner[0]) : tanh(inner[0]),
        abs(inner[1]) > 15.0 ? sign(inner[1]) : tanh(inner[1]),
        abs(inner[2]) > 15.0 ? sign(inner[2]) : tanh(inner[2]),
        abs(inner[3]) > 15.0 ? sign(inner[3]) : tanh(inner[3])
    );
    inner = 0.5 * (1.0 + inner);
    vec4 result = x * inner;
    return result;
`, q = `
    vec4 a2 = a * a;
    vec4 a3 = a2 * a;
    vec4 u  = ${S} * (a + ${w} * a3);
    vec4 t = vec4(
        abs(u[0]) > 15.0 ? sign(u[0]) : tanh(u[0]),
        abs(u[1]) > 15.0 ? sign(u[1]) : tanh(u[1]),
        abs(u[2]) > 15.0 ? sign(u[2]) : tanh(u[2]),
        abs(u[3]) > 15.0 ? sign(u[3]) : tanh(u[3])
    );
    vec4 sech2 = 1.0 - t * t;
    vec4 du_dx = ${S} * (1.0 + 3.0 * ${w} * a2);
    vec4 dgelu = 0.5 * (1.0 + t) + 0.5 * a * sech2 * du_dx;
    return dgelu * b;
`, se = 1e3;
function N({
  a: t,
  b: e,
  transposeA: s,
  transposeB: a,
  backend: n,
  activationSnippet: o,
  multiplier: r
}) {
  const i = t.shape.length, u = e.shape.length, l = s ? t.shape[i - 2] : t.shape[i - 1], p = a ? e.shape[u - 1] : e.shape[u - 2], h = s ? t.shape[i - 1] : t.shape[i - 2], d = a ? e.shape[u - 2] : e.shape[u - 1], x = t.shape.slice(0, -2), b = e.shape.slice(0, -2), m = G(x), c = G(b), $ = L(t.shape.slice(0, -2), e.shape.slice(0, -2)).concat([h, d]);
  F(
    l === p,
    () => `Error in matMul: inner shapes (${l}) and (${p}) of Tensors with shapes ${t.shape} and ${e.shape} and transposeA=${s} and transposeB=${a} must match.`
  );
  const f = s ? [m, l, h] : [m, h, l], v = a ? [c, d, p] : [c, p, d], A = M({ inputs: { x: t }, backend: n, attrs: { shape: f } }), y = M({ inputs: { x: e }, backend: n, attrs: { shape: v } }), D = [A, y], O = Math.max(m, c), _ = o, E = U(t.dtype, e.dtype), T = new W(
    f,
    v,
    [O, h, d],
    s,
    a,
    !1,
    _,
    !!r,
    !1
  ), k = [A, y];
  r && k.push(r);
  const z = n.runWebGLProgram(T, k, E), K = M({ inputs: { x: z }, backend: n, attrs: { shape: $ } });
  D.push(z);
  for (const P of D)
    n.disposeIntermediateTensorInfo(P);
  return K;
}
function J(t) {
  const { inputs: e, backend: s } = t, { x: a, kernel: n } = e;
  if (a === void 0 || n === void 0)
    throw new Error("BatchMatMul requires two input tensors.");
  return N({
    a,
    b: n,
    transposeA: !1,
    transposeB: !1,
    backend: s,
    activationSnippet: j
  });
}
const Q = {
  kernelName: "MatMulGelu",
  backendName: "webgl",
  kernelFunc: J
};
C(Q);
function V(t) {
  const { dy: e, x: s, kernel: a } = t.inputs, n = t.backend;
  return R(() => {
    const o = I().makeTensorFromTensorInfo(
      N({
        a: s,
        b: a,
        transposeA: !1,
        transposeB: !1,
        backend: n,
        activationSnippet: q,
        multiplier: e
      })
    ), r = B(o, a, !1, !0), i = B(s, o, !0, !1);
    return [r, i];
  });
}
const X = {
  kernelName: "MatMulGeluGrad",
  backendName: "webgl",
  kernelFunc: V
};
C(X);
export {
  W as M,
  se as a,
  N as b,
  J as c
};
