volvoxai 0.1.0

package/shaders/nonMaxSuppression.wgsl

@@ -0,0 +1,92 @@
+@group(0) @binding(0) var<storage, read> boxes : array<f32>;
+@group(0) @binding(1) var<storage, read> scores : array<f32>;
+@group(0) @binding(2) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    batches : u32,
+    spatial : u32,
+    classes : u32,
+    max_output : u32,
+    output_rows : u32,
+    iou_threshold : f32,
+    score_threshold : f32,
+}
+@group(0) @binding(3) var<uniform> params : Params;
+
+fn box_iou(b : u32, a_idx : u32, b_idx : u32) -> f32 {
+    let base_a = (b * params.spatial + a_idx) * 4u;
+    let base_b = (b * params.spatial + b_idx) * 4u;
+    let ay1 = boxes[base_a + 0u];
+    let ax1 = boxes[base_a + 1u];
+    let ay2 = boxes[base_a + 2u];
+    let ax2 = boxes[base_a + 3u];
+    let by1 = boxes[base_b + 0u];
+    let bx1 = boxes[base_b + 1u];
+    let by2 = boxes[base_b + 2u];
+    let bx2 = boxes[base_b + 3u];
+    let xx1 = max(ax1, bx1);
+    let yy1 = max(ay1, by1);
+    let xx2 = min(ax2, bx2);
+    let yy2 = min(ay2, by2);
+    let w = max(0.0, xx2 - xx1);
+    let h = max(0.0, yy2 - yy1);
+    let inter = w * h;
+    let area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1);
+    let area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1);
+    let denom = area_a + area_b - inter;
+    if (denom <= 0.0) { return 0.0; }
+    return inter / denom;
+}
+
+fn selected_suppresses(b : u32, c : u32, candidate : u32, out_start : u32, selected_count : u32) -> bool {
+    for (var j = 0u; j < selected_count; j = j + 1u) {
+        let row = out_start + j;
+        if (row >= params.output_rows) { return false; }
+        let prev_b = u32(output[row * 3u + 0u]);
+        let prev_c = u32(output[row * 3u + 1u]);
+        let prev_s = u32(output[row * 3u + 2u]);
+        if (prev_b == b && prev_c == c && box_iou(b, candidate, prev_s) > params.iou_threshold) {
+            return true;
+        }
+    }
+    return false;
+}
+
+@compute @workgroup_size(1)
+fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+    if (global_id.x != 0u) { return; }
+
+    for (var i = 0u; i < params.output_rows; i = i + 1u) {
+        output[i * 3u + 0u] = -1.0;
+        output[i * 3u + 1u] = -1.0;
+        output[i * 3u + 2u] = -1.0;
+    }
+
+    var out_idx = 0u;
+    for (var b = 0u; b < params.batches; b = b + 1u) {
+        for (var c = 0u; c < params.classes; c = c + 1u) {
+            let class_start = out_idx;
+            var selected = 0u;
+            loop {
+                if (selected >= params.max_output || out_idx >= params.output_rows) { break; }
+                var best_score = params.score_threshold;
+                var best_s = params.spatial;
+                for (var s = 0u; s < params.spatial; s = s + 1u) {
+                    let score = scores[b * (params.classes * params.spatial) + c * params.spatial + s];
+                    if (score < best_score) { continue; }
+                    if (selected_suppresses(b, c, s, class_start, selected)) { continue; }
+                    if (best_s == params.spatial || score > best_score) {
+                        best_score = score;
+                        best_s = s;
+                    }
+                }
+                if (best_s == params.spatial) { break; }
+                output[out_idx * 3u + 0u] = f32(b);
+                output[out_idx * 3u + 1u] = f32(c);
+                output[out_idx * 3u + 2u] = f32(best_s);
+                out_idx = out_idx + 1u;
+                selected = selected + 1u;
+            }
+        }
+    }
+}

package/shaders/pReLU.wgsl

@@ -0,0 +1,14 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read> weight : array<f32>;
+        @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32, c : u32 }
+        @group(0) @binding(3) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            let c_idx = idx % params.c;
+            let alpha = weight[c_idx];
+            let v = input[idx];
+            if (v > 0.0) { output[idx] = v; } else { output[idx] = v * alpha; }
+        }

package/shaders/pad.wgsl

@@ -0,0 +1,19 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { b: u32, in_h: u32, in_w: u32, c: u32, out_h: u32, out_w: u32, pt: u32, pl: u32, val: f32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            let total = params.b * params.out_h * params.out_w * params.c;
+            if (idx >= total) { return; }
+            let c = idx % params.c;
+            let x = (idx / params.c) % params.out_w;
+            let y = (idx / (params.c * params.out_w)) % params.out_h;
+            let b = idx / (params.c * params.out_w * params.out_h);
+            if (y >= params.pt && y < params.pt + params.in_h && x >= params.pl && x < params.pl + params.in_w) {
+                output[idx] = input[((b * params.in_h + (y - params.pt)) * params.in_w + (x - params.pl)) * params.c + c];
+            } else {
+                output[idx] = params.val;
+            }
+        }

package/shaders/profileX.wgsl

@@ -0,0 +1,28 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+            @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+            
+            struct Params { in_h : u32, in_w : u32, in_c : u32 }
+            @group(0) @binding(2) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let x = global_id.x;
+                let c = global_id.y;
+                
+                if (x >= params.in_w || c >= params.in_c) { return; }
+                
+                var max_val = -1e38;
+                var sum_val = 0.0;
+                
+                for (var y = 0u; y < params.in_h; y = y + 1u) {
+                    let val = input[(y * params.in_w + x) * params.in_c + c];
+                    if (val > max_val) { max_val = val; }
+                    sum_val = sum_val + val;
+                }
+                
+                let out_max_idx = c * params.in_w + x;
+                let out_mean_idx = (c + params.in_c) * params.in_w + x;
+                
+                output[out_max_idx] = max_val;
+                output[out_mean_idx] = sum_val / f32(params.in_h);
+            }

package/shaders/profileY.wgsl

@@ -0,0 +1,28 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+            @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+            
+            struct Params { in_h : u32, in_w : u32, in_c : u32 }
+            @group(0) @binding(2) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let y = global_id.x;
+                let c = global_id.y;
+                
+                if (y >= params.in_h || c >= params.in_c) { return; }
+                
+                var max_val = -1e38;
+                var sum_val = 0.0;
+                
+                for (var x = 0u; x < params.in_w; x = x + 1u) {
+                    let val = input[(y * params.in_w + x) * params.in_c + c];
+                    if (val > max_val) { max_val = val; }
+                    sum_val = sum_val + val;
+                }
+                
+                let out_max_idx = c * params.in_h + y;
+                let out_mean_idx = (c + params.in_c) * params.in_h + y;
+                
+                output[out_max_idx] = max_val;
+                output[out_mean_idx] = sum_val / f32(params.in_w);
+            }

package/shaders/quantizeLinear.wgsl

@@ -0,0 +1,69 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<u32>;
+
+struct Params {
+    size : u32,
+    input_zp : i32,
+    output_zp : i32,
+    has_input_scale : u32,
+    input_scale : f32,
+    output_scale : f32,
+    pad0 : u32,
+    pad1 : u32,
+}
+
+@group(0) @binding(2) var<uniform> params : Params;
+
+fn clamp_i8(v : i32) -> i32 {
+    return min(max(v, -128), 127);
+}
+
+fn i8_byte(v : i32) -> u32 {
+    if (v < 0) {
+        return u32(v + 256) & 255u;
+    }
+    return u32(v) & 255u;
+}
+
+fn round_even(v : f32) -> f32 {
+    let lo = floor(v);
+    let frac = v - lo;
+    if (frac < 0.5) {
+        return lo;
+    }
+    if (frac > 0.5) {
+        return lo + 1.0;
+    }
+    if (floor(lo * 0.5) == lo * 0.5) {
+        return lo;
+    }
+    return lo + 1.0;
+}
+
+fn quantize_one(idx : u32) -> u32 {
+    if (idx >= params.size) {
+        return 0u;
+    }
+    var v = input[idx];
+    if (params.has_input_scale == 1u) {
+        v = (v - f32(params.input_zp)) * params.input_scale;
+    }
+    let qf = round_even(v / params.output_scale + f32(params.output_zp));
+    let q = clamp_i8(i32(min(max(qf, -128.0), 127.0)));
+    return i8_byte(q);
+}
+
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+    let word_idx = global_id.x;
+    let base = word_idx * 4u;
+    if (base >= params.size) {
+        return;
+    }
+
+    var packed = quantize_one(base);
+    packed = packed | (quantize_one(base + 1u) << 8u);
+    packed = packed | (quantize_one(base + 2u) << 16u);
+    packed = packed | (quantize_one(base + 3u) << 24u);
+    output[word_idx] = packed;
+}

package/shaders/rMSNorm.wgsl

@@ -0,0 +1,21 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read> weight : array<f32>;
+        @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+        struct Params { seq_len : u32, d_model : u32, eps : f32 }
+        @group(0) @binding(3) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let s = global_id.x;
+            if (s >= params.seq_len) { return; }
+            var var_val = 0.0;
+            let offset = s * params.d_model;
+            for (var d = 0u; d < params.d_model; d = d + 1u) {
+                let v = input[offset + d];
+                var_val = var_val + v * v;
+            }
+            var_val = var_val / f32(params.d_model);
+            let inv_std = 1.0 / sqrt(var_val + params.eps);
+            for (var d = 0u; d < params.d_model; d = d + 1u) {
+                output[offset + d] = input[offset + d] * inv_std * weight[d];
+            }
+        }

package/shaders/reLU.wgsl

@@ -0,0 +1,13 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            let x = input[idx];
+            var out_val = x;
+            if (x > 0.0) { out_val = x; } else { out_val = 0.0; }
+            output[idx] = out_val;
+        }

package/shaders/reduce.wgsl

@@ -0,0 +1,17 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<f32>;
+// Reduce over the last (innermost) axis: `b` rows of length `d` -> `b` outputs.
+// `inv` scales the sum: 1.0 for ReduceSum, 1.0/d for ReduceMean.
+struct Params { b : u32, d : u32, inv : f32 }
+@group(0) @binding(2) var<uniform> p : Params;
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+  let i = gid.x;
+  if (i >= p.b) { return; }
+  let offset = i * p.d;
+  var sum = 0.0;
+  for (var j = 0u; j < p.d; j = j + 1u) {
+    sum = sum + input[offset + j];
+  }
+  output[i] = sum * p.inv;
+}

package/shaders/resize.wgsl

@@ -0,0 +1,52 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    n : u32,
+    h : u32,
+    w : u32,
+    c : u32,
+    out_h : u32,
+    out_w : u32,
+    mode : u32,
+}
+@group(0) @binding(2) var<uniform> params : Params;
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let z = gid.z;
+    let nb = z / params.c;
+    let ch = z - nb * params.c;
+    if (nb >= params.n || ox >= params.out_w || oy >= params.out_h) { return; }
+
+    if (params.mode == 0u) {
+        let iy = (oy * params.h) / params.out_h;
+        let ix = (ox * params.w) / params.out_w;
+        output[((nb * params.out_h + oy) * params.out_w + ox) * params.c + ch] =
+            input[((nb * params.h + iy) * params.w + ix) * params.c + ch];
+        return;
+    }
+
+    let scale_y = f32(params.h) / f32(params.out_h);
+    let scale_x = f32(params.w) / f32(params.out_w);
+    var fy = (f32(oy) + 0.5) * scale_y - 0.5;
+    var fx = (f32(ox) + 0.5) * scale_x - 0.5;
+    if (fy < 0.0) { fy = 0.0; }
+    if (fx < 0.0) { fx = 0.0; }
+    let y0 = u32(fy);
+    let x0 = u32(fx);
+    let y1 = min(y0 + 1u, params.h - 1u);
+    let x1 = min(x0 + 1u, params.w - 1u);
+    let dy = fy - f32(y0);
+    let dx = fx - f32(x0);
+    let base = nb * params.h * params.w * params.c;
+    let v00 = input[base + (y0 * params.w + x0) * params.c + ch];
+    let v01 = input[base + (y0 * params.w + x1) * params.c + ch];
+    let v10 = input[base + (y1 * params.w + x0) * params.c + ch];
+    let v11 = input[base + (y1 * params.w + x1) * params.c + ch];
+    let val = v00 * (1.0 - dy) * (1.0 - dx) + v01 * (1.0 - dy) * dx +
+              v10 * dy * (1.0 - dx) + v11 * dy * dx;
+    output[((nb * params.out_h + oy) * params.out_w + ox) * params.c + ch] = val;
+}

package/shaders/sDPA.wgsl

@@ -0,0 +1,71 @@
+@group(0) @binding(0) var<storage, read> qkv : array<f32>;
+            @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+            
+            struct Params {
+                seq_len : u32,
+                d_model : u32,
+                num_heads : u32,
+                head_dim : u32,
+                scale : f32,
+            }
+            @group(0) @binding(2) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64, 1, 1)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let q_idx = global_id.x; 
+                let h_idx = global_id.y; 
+                
+                if (q_idx >= params.seq_len || h_idx >= params.num_heads) { return; }
+                
+                let head_dim = params.head_dim;
+                let d_model = params.d_model;
+                
+                var max_logit : f32 = -1e38;
+                
+                // Cache Q for this head and this q_idx
+                var q_cache : array<f32, 64>; // max head_dim 64
+                for (var d = 0u; d < head_dim; d = d + 1u) {
+                    q_cache[d] = qkv[q_idx * (d_model * 3u) + (h_idx * head_dim) + d];
+                }
+                
+                // Pass 1: find max
+                for (var k_idx = 0u; k_idx <= q_idx; k_idx = k_idx + 1u) {
+                    var score : f32 = 0.0;
+                    for (var d = 0u; d < head_dim; d = d + 1u) {
+                        let k_val = qkv[k_idx * (d_model * 3u) + d_model + (h_idx * head_dim) + d];
+                        score = score + (q_cache[d] * k_val);
+                    }
+                    score = score * params.scale;
+                    if (score > max_logit) { max_logit = score; }
+                }
+                
+                // Pass 2: sum exp
+                var sum_exp : f32 = 0.0;
+                for (var k_idx = 0u; k_idx <= q_idx; k_idx = k_idx + 1u) {
+                    var score : f32 = 0.0;
+                    for (var d = 0u; d < head_dim; d = d + 1u) {
+                        let k_val = qkv[k_idx * (d_model * 3u) + d_model + (h_idx * head_dim) + d];
+                        score = score + (q_cache[d] * k_val);
+                    }
+                    score = score * params.scale;
+                    sum_exp = sum_exp + exp(score - max_logit);
+                }
+                
+                // Pass 3: output
+                for (var d = 0u; d < head_dim; d = d + 1u) {
+                    var out_val : f32 = 0.0;
+                    for (var k_idx = 0u; k_idx <= q_idx; k_idx = k_idx + 1u) {
+                        var score : f32 = 0.0;
+                        for (var kd = 0u; kd < head_dim; kd = kd + 1u) {
+                            let k_val = qkv[k_idx * (d_model * 3u) + d_model + (h_idx * head_dim) + kd];
+                            score = score + (q_cache[kd] * k_val);
+                        }
+                        score = score * params.scale;
+                        
+                        let w = exp(score - max_logit) / sum_exp;
+                        let v_val = qkv[k_idx * (d_model * 3u) + d_model * 2u + (h_idx * head_dim) + d];
+                        out_val = out_val + (w * v_val);
+                    }
+                    output[q_idx * d_model + (h_idx * head_dim) + d] = out_val;
+                }
+            }

package/shaders/siLU.wgsl

@@ -0,0 +1,13 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            let x = input[idx];
+            var out_val = x;
+            out_val = x * (1.0 / (1.0 + exp(-x)));
+            output[idx] = out_val;
+        }

package/shaders/sigmoid.wgsl

@@ -0,0 +1,13 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            let x = input[idx];
+            var out_val = x;
+            out_val = 1.0 / (1.0 + exp(-x));
+            output[idx] = out_val;
+        }

package/shaders/slice.wgsl

@@ -0,0 +1,26 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<f32>;
+// General strided slice over 4 dims (leading dims are padded to 1). For each
+// output element the source coord along axis k is start[k] + out_coord * step[k].
+struct Params {
+  out_b: u32, out_c: u32, out_h: u32, out_w: u32,
+  in_c: u32, in_h: u32, in_w: u32,
+  s0: u32, s1: u32, s2: u32, s3: u32,
+  st0: u32, st1: u32, st2: u32, st3: u32,
+  total: u32
+}
+@group(0) @binding(2) var<uniform> p : Params;
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+  let idx = gid.x;
+  if (idx >= p.total) { return; }
+  let ow = idx % p.out_w;
+  let oh = (idx / p.out_w) % p.out_h;
+  let oc = (idx / (p.out_w * p.out_h)) % p.out_c;
+  let ob = idx / (p.out_w * p.out_h * p.out_c);
+  let ib = p.s0 + ob * p.st0;
+  let ic = p.s1 + oc * p.st1;
+  let ih = p.s2 + oh * p.st2;
+  let iw = p.s3 + ow * p.st3;
+  output[idx] = input[((ib * p.in_c + ic) * p.in_h + ih) * p.in_w + iw];
+}

package/shaders/softmax.wgsl

@@ -0,0 +1,23 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { b : u32, d : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let i = global_id.x;
+            if (i >= params.b) { return; }
+            let offset = i * params.d;
+            var max_val = -100000.0;
+            for (var j = 0u; j < params.d; j = j + 1u) {
+                if (input[offset + j] > max_val) { max_val = input[offset + j]; }
+            }
+            var sum = 0.0;
+            for (var j = 0u; j < params.d; j = j + 1u) {
+                let e = exp(input[offset + j] - max_val);
+                output[offset + j] = e;
+                sum = sum + e;
+            }
+            for (var j = 0u; j < params.d; j = j + 1u) {
+                output[offset + j] = output[offset + j] / sum;
+            }
+        }

package/shaders/spatialSoftargmaxY.wgsl

@@ -0,0 +1,32 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+            @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+            
+            struct Params { in_h : u32, in_w : u32, in_c : u32 }
+            @group(0) @binding(2) var<uniform> params : Params;
+            
+            @compute @workgroup_size(64)
+            fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+                let x = global_id.x;
+                let c = global_id.y;
+                
+                if (x >= params.in_w || c >= params.in_c) { return; }
+                
+                var max_val = -1e38;
+                for (var y = 0u; y < params.in_h; y = y + 1u) {
+                    let val = input[(y * params.in_w + x) * params.in_c + c];
+                    if (val > max_val) { max_val = val; }
+                }
+                
+                var denom = 0.0;
+                var weighted = 0.0;
+                for (var y = 0u; y < params.in_h; y = y + 1u) {
+                    let val = input[(y * params.in_w + x) * params.in_c + c];
+                    let ev = exp(val - max_val);
+                    denom = denom + ev;
+                    weighted = weighted + ev * (f32(y) + 0.5) / f32(params.in_h);
+                }
+                
+                var out_val = 0.0;
+                if (denom > 0.0) { out_val = weighted / denom; }
+                output[c * params.in_w + x] = out_val;
+            }

package/shaders/split.wgsl

@@ -0,0 +1,15 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { total : u32, inner : u32, split_size : u32, axis_in : u32, offset : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+            let i = gid.x;
+            if (i >= params.total) { return; }
+            let inner_idx = i % params.inner;
+            let s = (i / params.inner) % params.split_size;
+            let outer_idx = i / (params.split_size * params.inner);
+            let in_idx = outer_idx * (params.axis_in * params.inner)
+                       + (params.offset + s) * params.inner + inner_idx;
+            output[i] = input[in_idx];
+        }

package/shaders/sub.wgsl

@@ -0,0 +1,34 @@
+@group(0) @binding(0) var<storage, read> a : array<f32>;
+        @group(0) @binding(1) var<storage, read> b : array<f32>;
+        @group(0) @binding(2) var<storage, read_write> output : array<f32>;
+        
+        struct Params { size : u32, is_b_scalar : u32, b_size : u32, a_size: u32, is_a_scalar: u32 }
+        @group(0) @binding(3) var<uniform> params : Params;
+        
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            
+            var a_val : f32 = 0.0;
+            if (params.is_a_scalar == 1u) {
+                a_val = a[0];
+            } else if (params.a_size < params.size && params.a_size > 0u) {
+                a_val = a[idx % params.a_size];
+            } else {
+                a_val = a[idx];
+            }
+            
+            var b_val : f32 = 0.0;
+            if (params.is_b_scalar == 1u) {
+                b_val = b[0];
+            } else if (params.b_size < params.size && params.b_size > 0u) {
+                b_val = b[idx % params.b_size];
+            } else {
+                b_val = b[idx];
+            }
+            
+            var out_val : f32 = 0.0;
+            out_val = a_val - b_val;
+            output[idx] = out_val;
+        }

package/shaders/tanh.wgsl

@@ -0,0 +1,13 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+        @group(0) @binding(1) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32 }
+        @group(0) @binding(2) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            let x = input[idx];
+            var out_val = x;
+            let e2x = exp(2.0 * x); out_val = (e2x - 1.0) / (e2x + 1.0);
+            output[idx] = out_val;
+        }

package/shaders/upsample2x.wgsl

@@ -0,0 +1,24 @@
+@group(0) @binding(0) var<storage, read> input : array<f32>;
+@group(0) @binding(1) var<storage, read_write> output : array<f32>;
+
+struct Params {
+    n : u32,
+    h : u32,
+    w : u32,
+    c : u32,
+}
+@group(0) @binding(2) var<uniform> params : Params;
+
+@compute @workgroup_size(8, 8, 1)
+fn main(@builtin(global_invocation_id) gid : vec3<u32>) {
+    let ox = gid.x;
+    let oy = gid.y;
+    let z = gid.z;
+    let nb = z / params.c;
+    let ch = z - nb * params.c;
+    let out_h = params.h * 2u;
+    let out_w = params.w * 2u;
+    if (nb >= params.n || ox >= out_w || oy >= out_h) { return; }
+    output[((nb * out_h + oy) * out_w + ox) * params.c + ch] =
+        input[((nb * params.h + (oy / 2u)) * params.w + (ox / 2u)) * params.c + ch];
+}

package/shaders/where.wgsl

@@ -0,0 +1,12 @@
+@group(0) @binding(0) var<storage, read> cond : array<f32>;
+        @group(0) @binding(1) var<storage, read> a : array<f32>;
+        @group(0) @binding(2) var<storage, read> b : array<f32>;
+        @group(0) @binding(3) var<storage, read_write> output : array<f32>;
+        struct Params { size : u32 }
+        @group(0) @binding(4) var<uniform> params : Params;
+        @compute @workgroup_size(64)
+        fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
+            let idx = global_id.x;
+            if (idx >= params.size) { return; }
+            if (cond[idx] != 0.0) { output[idx] = a[idx]; } else { output[idx] = b[idx]; }
+        }