import { a as h } from "../../index-DSGwv2Yx.js";
class g {
  variableNames = ["x", "sin", "cos"];
  outputShape;
  userCode;
  // enableShapeUniforms = true;
  customUniforms = [{ name: "pastLen", type: "int" }];
  constructor(t, s, e, n) {
    this.outputShape = [t, s, e, n], this.userCode = `
        void main() {
            ivec4 coords = getOutputCoords(); // [b, h, t, d]
            int b = coords.x;
            int h = coords.y;
            int t = coords.z;
            int d = coords.w;

            int rotaryDim = ${n};

            float outVal = 0.0;

            if (d < rotaryDim) {
                int pairIdx = d / 2;
                float cos = getCos(t + pastLen, pairIdx, 0);
                float sin = getSin(t + pastLen, pairIdx, 0);

                if (d % 2 == 0) {
                    // even index
                    float even = getX(b, h, t, d);
                    float odd = getX(b, h, t, d + 1);
                    outVal = even * cos - odd * sin;
                } else {
                    // odd index
                    float even = getX(b, h, t, d - 1);
                    float odd = getX(b, h, t, d);
                    outVal = even * sin + odd * cos;
                }
            } else {
                // pass through for non-rotary dims
                outVal = getX(b, h, t, d);
            }

            setOutput(outVal);
        }
        `;
  }
}
function f(o) {
  const { x: t } = o.inputs, { pastLen: s, ropeCache: e, negSin: n } = o.attrs, a = n ? e.getNegSin() : e.getSin(), r = e.getCos(), d = o.backend, i = t.shape[0], c = t.shape[1], p = t.shape[2], u = t.shape[3], l = new g(i, c, p, u);
  return d.runWebGLProgram(l, [t, a, r], "float32", [[s]]);
}
const m = {
  kernelName: "Rope",
  backendName: "webgl",
  kernelFunc: f
};
h(m);
