import { e as u } from "../../webgpu_program-CbjdYLYk.js";
import { f as d, c as i } from "../../webgpu_util-DuofJBMo.js";
import { a as p } from "../../index-DSGwv2Yx.js";
class c {
  variableNames = ["x"];
  outputShape;
  shaderKey = "Dropout16";
  dispatchLayout;
  dispatch;
  workgroupSize = [64, 1, 1];
  size = !0;
  uniforms = "dropout: f32, seed: f32";
  constructor(e) {
    this.shaderKey = "Dropout16", this.outputShape = e, this.dispatchLayout = d(this.outputShape), this.dispatch = i(this.dispatchLayout, this.outputShape, this.workgroupSize);
  }
  getUserCode() {
    return `

        fn random(coords: vec3<i32>) -> vec2<f32> {
            let x1 = f32(coords.x * 4096 + coords.y * 256 + coords.z * 2 * 16);
            let x2 = f32(coords.x * 4096 + coords.y * 256 + (coords.z * 2 + 1) * 16);
            return vec2<f32>(fract(sin(uniforms.seed + x1) * 43758.5453123), fract(sin(uniforms.seed + x2) * 43758.5453123));
        }

        ${u("index")} {
            if (index < uniforms.size) {
                let coords = getCoordsFromIndex(index);
                let values = unpack2x16float(u32(x[index]));
                let keepProb = 1.0 - uniforms.dropout;
                let rand = random(coords);
                let mask = step(rand, vec2<f32>(keepProb));
                let outVal = values * mask / keepProb;
                result[index] = i32(pack2x16float(outVal));
            }
        }
        `;
  }
}
function f(t) {
  const { x: e } = t.inputs, { dropout: o, seed: r } = t.attrs, s = t.backend, a = new c(e.shape), n = [
    { type: "float32", data: [o] },
    { type: "float32", data: [r] }
  ];
  return s.runWebGPUProgram(a, [e], "packedF16", n);
}
const l = {
  kernelName: "Dropout16",
  backendName: "webgpu",
  kernelFunc: f
};
p(l);
