//! Runtime codegen for K-specialized Taylor forward GPU shaders.
//!
//! Generates fully unrolled WGSL and CUDA shader source for Taylor jets of order
//! K ∈ {1, 2, 3, 4, 5}. Each generated shader has K-specific Cauchy products and
//! recurrences with no dynamic loops, preserving the performance of the handwritten
//! K=3 shader while supporting arbitrary orders.

use std::fmt::Write;

/// Generate a complete WGSL shader for K-th order Taylor forward propagation.
///
/// The generated shader uses the standard Taylor bind group layout
/// with K coefficients per jet (K=1..5).
///
/// # Panics
/// Panics if `k` is not in 1..=5.
#[must_use]
pub fn generate_taylor_wgsl(k: usize) -> String {
    assert!((1..=5).contains(&k), "K must be in 1..=5, got {k}");

    let mut s = String::with_capacity(16384);
    writeln!(s, "// Auto-generated Taylor forward K={k} kernel.").unwrap();
    writeln!(s, "// Do not edit — generated by taylor_codegen.rs.\n").unwrap();

    // Opcode constants
    write_wgsl_opcodes(&mut s);

    // TapeMeta, bind groups
    write_wgsl_bindings(&mut s, k);

    // Helper builtins missing from WGSL
    write_wgsl_helpers(&mut s);

    // JetK struct and operations
    write_wgsl_jet_type(&mut s, k);
    write_wgsl_jet_arithmetic(&mut s, k);
    write_wgsl_jet_transcendental(&mut s, k);
    write_wgsl_jet_inverse_trig(&mut s, k);

    // Main kernel
    write_wgsl_main_kernel(&mut s, k);

    s
}

/// Generate a complete CUDA kernel for K-th order Taylor forward propagation.
///
/// Expects `#define FLOAT_TYPE float` or `double` to be prepended by the caller.
///
/// # Panics
/// Panics if `k` is not in 1..=5.
#[must_use]
pub fn generate_taylor_cuda(k: usize) -> String {
    assert!((1..=5).contains(&k), "K must be in 1..=5, got {k}");

    let mut s = String::with_capacity(16384);
    writeln!(s, "// Auto-generated Taylor forward K={k} kernel.").unwrap();
    writeln!(s, "// Do not edit — generated by taylor_codegen.rs.\n").unwrap();
    writeln!(s, "typedef FLOAT_TYPE F;").unwrap();

    // Opcode constants
    write_cuda_opcodes(&mut s);

    // Math helpers
    write_cuda_helpers(&mut s);

    // JetK struct and operations
    write_cuda_jet_type(&mut s, k);
    write_cuda_jet_arithmetic(&mut s, k);
    write_cuda_jet_transcendental(&mut s, k);
    write_cuda_jet_inverse_trig(&mut s, k);

    // Main kernel
    write_cuda_main_kernel(&mut s, k);

    s
}

// ══════════════════════════════════════════════
//  WGSL code generation
// ══════════════════════════════════════════════

fn write_wgsl_opcodes(s: &mut String) {
    let ops = [
        ("OP_INPUT", 0),
        ("OP_CONST", 1),
        ("OP_ADD", 2),
        ("OP_SUB", 3),
        ("OP_MUL", 4),
        ("OP_DIV", 5),
        ("OP_REM", 6),
        ("OP_POWF", 7),
        ("OP_ATAN2", 8),
        ("OP_HYPOT", 9),
        ("OP_MAX", 10),
        ("OP_MIN", 11),
        ("OP_NEG", 12),
        ("OP_RECIP", 13),
        ("OP_SQRT", 14),
        ("OP_CBRT", 15),
        ("OP_POWI", 16),
        ("OP_EXP", 17),
        ("OP_EXP2", 18),
        ("OP_EXPM1", 19),
        ("OP_LN", 20),
        ("OP_LOG2", 21),
        ("OP_LOG10", 22),
        ("OP_LN1P", 23),
        ("OP_SIN", 24),
        ("OP_COS", 25),
        ("OP_TAN", 26),
        ("OP_ASIN", 27),
        ("OP_ACOS", 28),
        ("OP_ATAN", 29),
        ("OP_SINH", 30),
        ("OP_COSH", 31),
        ("OP_TANH", 32),
        ("OP_ASINH", 33),
        ("OP_ACOSH", 34),
        ("OP_ATANH", 35),
        ("OP_ABS", 36),
        ("OP_SIGNUM", 37),
        ("OP_FLOOR", 38),
        ("OP_CEIL", 39),
        ("OP_ROUND", 40),
        ("OP_TRUNC", 41),
        ("OP_FRACT", 42),
    ];
    for (name, val) in &ops {
        writeln!(s, "const {name}: u32 = {val}u;").unwrap();
    }
    writeln!(s).unwrap();
}

fn write_wgsl_bindings(s: &mut String, k: usize) {
    writeln!(
        s,
        "struct TapeMeta {{
    num_ops: u32,
    num_inputs: u32,
    num_variables: u32,
    num_outputs: u32,
    batch_size: u32,
    _pad0: u32,
    _pad1: u32,
    _pad2: u32,
}}

@group(0) @binding(0) var<storage, read> opcodes: array<u32>;
@group(0) @binding(1) var<storage, read> arg0: array<u32>;
@group(0) @binding(2) var<storage, read> arg1: array<u32>;
@group(0) @binding(3) var<storage, read> constants: array<f32>;
@group(0) @binding(4) var<uniform> tape_meta: TapeMeta;
@group(0) @binding(5) var<storage, read> output_indices: array<u32>;

@group(1) @binding(0) var<storage, read> primal_inputs: array<f32>;
@group(1) @binding(1) var<storage, read> direction_seeds: array<f32>;
@group(1) @binding(2) var<storage, read_write> jets: array<f32>;
@group(1) @binding(3) var<storage, read_write> jet_outputs: array<f32>;
"
    )
    .unwrap();
    let _ = k; // K only affects buffer sizes at dispatch time, not bindings
}

fn write_wgsl_helpers(s: &mut String) {
    writeln!(
        s,
        "fn sinh_f(x: f32) -> f32 {{ return (exp(x) - exp(-x)) * 0.5; }}
fn cosh_f(x: f32) -> f32 {{ return (exp(x) + exp(-x)) * 0.5; }}
fn asinh_f(x: f32) -> f32 {{ let a = abs(x); if a > 1e9 {{ let r = log(a) + log(1.0 + sqrt(1.0 + 1.0 / (a * a))); return select(-r, r, x >= 0.0); }} let r = log(a + sqrt(a * a + 1.0)); return select(-r, r, x >= 0.0); }}
fn acosh_f(x: f32) -> f32 {{ if x > 1e9 {{ return log(x) + log(1.0 + sqrt(1.0 - 1.0 / (x * x))); }} return log(x + sqrt((x - 1.0) * (x + 1.0))); }}
fn atanh_f(x: f32) -> f32 {{ return 0.5 * log((1.0 + x) / (1.0 - x)); }}
fn powf_real(base: f32, b: f32) -> f32 {{
    // WGSL `pow(x, y)` is undefined for x < 0. Rust/C `powf` define x^y for
    // x < 0 only at integer y: sign(x)^y * |x|^y; non-integer y is NaN.
    if base >= 0.0 {{ return pow(base, b); }}
    let rb = round(b);
    if rb != b {{ return bitcast<f32>(0x7fc00000u); }}
    let mag = pow(abs(base), b);
    if (i32(rb) & 1) != 0 {{ return -mag; }}
    return mag;
}}
fn expm1_f(x: f32) -> f32 {{
    // Accurate exp(x)-1. For small |x| the naive `exp(x)-1` cancels catastrophically
    // (exp(x) ~ 1), so use the Taylor series (no cancellation). For |x| >= 0.125 the
    // naive form is accurate and gives the right limits directly (+Inf for large x,
    // -1 for very negative x). The Kahan `(u-1)*x/log(u)` trick is avoided: Metal
    // fast-math both reassociates its grouping (overflowing `um1*x` for large x where
    // expm1 is finite) and, for log1p, folds `(1+x)-1` back to `x`.
    if (abs(x) < 0.125) {{
        // x + x^2/2 + x^3/6 + x^4/24 + x^5/120 + x^6/720  (Horner)
        return x * (1.0 + x * (0.5 + x * (0.16666667 + x * (0.041666668 + x * (0.008333334 + x * 0.0013888889)))));
    }}
    return exp(x) - 1.0;
}}
fn is_nan_f32(x: f32) -> bool {{
    // `x != x` is unreliable under Metal fast-math (it may assume no NaNs); test the
    // IEEE exponent/mantissa bits directly: NaN is exponent all-ones + non-zero mantissa,
    // i.e. `(bits & 0x7fffffff) > 0x7f800000` (0x7f800000 is exactly +Inf).
    return (bitcast<u32>(x) & 0x7fffffffu) > 0x7f800000u;
}}
fn signum_f32(x: f32) -> f32 {{
    // Rust f32::signum: -1 for -0.0 (sign bit set), +1 for +0.0/positive, NaN at NaN.
    // `x >= 0.0` wrongly maps -0.0 to +1; inspect the sign bit instead.
    if (is_nan_f32(x)) {{ return x; }}
    return select(1.0, -1.0, (bitcast<u32>(x) & 0x80000000u) != 0u);
}}
fn log1p_f(x: f32) -> f32 {{
    // Accurate log(1+x). For small |x| the naive `log(1 + x)` loses the low bits
    // of `x` when it is added to 1.0; the Kahan `log(u)*x/(u-1)` trick is defeated
    // by Metal fast-math folding `(1+x)-1` back to `x`, so use the Mercator series
    // (no cancellation) below a threshold and the direct log elsewhere (where
    // `1+x` retains enough bits). Only affects the primal v[0]; higher coeffs use
    // the jet_ln recurrence. x=-1 (>= threshold in magnitude) flows to log(0)=-Inf.
    if (abs(x) < 0.125) {{
        // x - x^2/2 + x^3/3 - x^4/4 + x^5/5 - x^6/6  (Horner)
        let p = 1.0 + x * (-0.5 + x * (0.33333334 + x * (-0.25 + x * (0.2 + x * (-0.16666667)))));
        return x * p;
    }}
    return log(1.0 + x);
}}
"
    )
    .unwrap();
}

/// Emit a WGSL domain guard at the top of a jet fn: if `cond` holds (leading
/// coeff strictly outside the real domain), return an all-NaN jet — mirroring the
/// CPU `taylor_*` guards, which short-circuit the whole jet rather than leave a
/// NaN primal beside finite higher coefficients.
fn write_wgsl_jet_nan_guard(s: &mut String, k: usize, cond: &str) {
    writeln!(s, "    if ({cond}) {{").unwrap();
    writeln!(s, "        var n: JetK;").unwrap();
    for i in 0..k {
        writeln!(s, "        n.v[{i}] = bitcast<f32>(0x7fc00000u);").unwrap();
    }
    writeln!(s, "        return n;").unwrap();
    writeln!(s, "    }}").unwrap();
}

/// CUDA counterpart of [`write_wgsl_jet_nan_guard`]; uses the emitter's `F(0.0/0.0)`
/// NaN idiom.
fn write_cuda_jet_nan_guard(s: &mut String, k: usize, cond: &str) {
    writeln!(s, "    if ({cond}) {{").unwrap();
    writeln!(s, "        JetK n;").unwrap();
    for i in 0..k {
        writeln!(s, "        n.v[{i}] = F(0.0/0.0);").unwrap();
    }
    writeln!(s, "        return n;").unwrap();
    writeln!(s, "    }}").unwrap();
}

fn write_wgsl_jet_type(s: &mut String, k: usize) {
    writeln!(s, "struct JetK {{ v: array<f32, {k}>, }}\n").unwrap();

    // jet_const: create jet from scalar
    write!(
        s,
        "fn jet_const(val: f32) -> JetK {{\n    var j: JetK;\n    j.v[0] = val;\n"
    )
    .unwrap();
    for i in 1..k {
        writeln!(s, "    j.v[{i}] = 0.0;").unwrap();
    }
    writeln!(s, "    return j;\n}}\n").unwrap();

    // jet_load: read jet from buffer
    writeln!(s, "fn jet_load(base: u32) -> JetK {{").unwrap();
    writeln!(s, "    var j: JetK;").unwrap();
    for i in 0..k {
        writeln!(s, "    j.v[{i}] = jets[base + {i}u];").unwrap();
    }
    writeln!(s, "    return j;\n}}\n").unwrap();

    // jet_store: write jet to buffer
    writeln!(s, "fn jet_store(base: u32, j: JetK) {{").unwrap();
    for i in 0..k {
        writeln!(s, "    jets[base + {i}u] = j.v[{i}];").unwrap();
    }
    writeln!(s, "}}\n").unwrap();
}

fn write_wgsl_jet_arithmetic(s: &mut String, k: usize) {
    // Add
    writeln!(s, "fn jet_add(a: JetK, b: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    for i in 0..k {
        writeln!(s, "    c.v[{i}] = a.v[{i}] + b.v[{i}];").unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Sub
    writeln!(s, "fn jet_sub(a: JetK, b: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    for i in 0..k {
        writeln!(s, "    c.v[{i}] = a.v[{i}] - b.v[{i}];").unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Neg
    writeln!(s, "fn jet_neg(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    for i in 0..k {
        writeln!(s, "    c.v[{i}] = -a.v[{i}];").unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Scale
    writeln!(s, "fn jet_scale(a: JetK, s: f32) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    for i in 0..k {
        writeln!(s, "    c.v[{i}] = a.v[{i}] * s;").unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Mul (Cauchy product): c[i] = Σ_{j=0}^{i} a[j] * b[i-j]
    writeln!(s, "fn jet_mul(a: JetK, b: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    for i in 0..k {
        let mut terms = Vec::new();
        for j in 0..=i {
            terms.push(format!("a.v[{j}] * b.v[{}]", i - j));
        }
        writeln!(s, "    c.v[{i}] = {};", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Div: c[i] = (a[i] - Σ_{j=1}^{i} b[j]*c[i-j]) / b[0]
    writeln!(s, "fn jet_div(a: JetK, b: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    let inv_b0 = 1.0 / b.v[0];").unwrap();
    writeln!(s, "    c.v[0] = a.v[0] * inv_b0;").unwrap();
    for i in 1..k {
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("b.v[{j}] * c.v[{}]", i - j));
        }
        writeln!(
            s,
            "    c.v[{i}] = (a.v[{i}] - ({})) * inv_b0;",
            terms.join(" + ")
        )
        .unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Recip: c = 1/a — special case: c[i] = -(Σ_{j=1}^{i} a[j]*c[i-j]) * c[0]
    writeln!(s, "fn jet_recip(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = 1.0 / a.v[0];").unwrap();
    for i in 1..k {
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("a.v[{j}] * c.v[{}]", i - j));
        }
        writeln!(s, "    c.v[{i}] = -({}) * c.v[0];", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();
}

fn write_wgsl_jet_transcendental(s: &mut String, k: usize) {
    // Exp: c[0] = exp(a[0]), c[i] = (1/i) * Σ_{j=1}^{i} j*a[j]*c[i-j]
    writeln!(s, "fn jet_exp(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = exp(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * c.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = {inv_i:.10} * ({});", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Ln: c[0] = ln(a[0]), c[i] = (a[i] - (1/i)*Σ_{j=1}^{i-1} j*c[j]*a[i-j]) / a[0]
    writeln!(s, "fn jet_ln(a: JetK) -> JetK {{").unwrap();
    // Domain guard: ln(a0<0) is NaN, but the higher coeffs (built from 1/a0) stay
    // finite — short-circuit to an all-NaN jet (mirrors CPU taylor_ln). Strict
    // `< 0` keeps a0==0 as the IEEE -Inf singularity. log2/log10/ln1p delegate to
    // jet_ln, so this one guard covers them too.
    write_wgsl_jet_nan_guard(s, k, "a.v[0] < 0.0");
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = log(a.v[0]);").unwrap();
    // `inv_a0` and the recurrence only exist for K > 1; emitting the `let` at
    // K=1 (where the loop below is empty) leaves an unused binding in the shader.
    if k > 1 {
        writeln!(s, "    let inv_a0 = 1.0 / a.v[0];").unwrap();
        for i in 1..k {
            if i == 1 {
                writeln!(s, "    c.v[1] = a.v[1] * inv_a0;").unwrap();
            } else {
                let inv_i = 1.0 / i as f64;
                let mut terms = Vec::new();
                for j in 1..i {
                    terms.push(format!("{:.1} * c.v[{j}] * a.v[{}]", j as f64, i - j));
                }
                writeln!(
                    s,
                    "    c.v[{i}] = (a.v[{i}] - {inv_i:.10} * ({})) * inv_a0;",
                    terms.join(" + ")
                )
                .unwrap();
            }
        }
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Sqrt: c[0] = sqrt(a[0]), c[i] = (a[i] - Σ_{j=1}^{i-1} c[j]*c[i-j]) / (2*c[0])
    writeln!(s, "fn jet_sqrt(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    // sqrt(0)=0 has a vertical tangent: the higher coeffs are uniformly +Inf, not
    // the sign-dependent `a.v[i]*(0.5/0)` the recurrence would give (matches CPU
    // taylor_sqrt). a.v[0] < 0 already all-NaNs via sqrt(neg) in the recurrence.
    writeln!(s, "    if (a.v[0] == 0.0) {{").unwrap();
    writeln!(s, "        c.v[0] = 0.0;").unwrap();
    for i in 1..k {
        writeln!(s, "        c.v[{i}] = bitcast<f32>(0x7f800000u);").unwrap();
    }
    writeln!(s, "        return c;").unwrap();
    writeln!(s, "    }}").unwrap();
    writeln!(s, "    c.v[0] = sqrt(a.v[0]);").unwrap();
    if k > 1 {
        writeln!(s, "    let inv_2c0 = 0.5 / c.v[0];").unwrap();
    }
    for i in 1..k {
        if i == 1 {
            writeln!(s, "    c.v[1] = a.v[1] * inv_2c0;").unwrap();
        } else {
            let mut terms = Vec::new();
            for j in 1..i {
                terms.push(format!("c.v[{j}] * c.v[{}]", i - j));
            }
            writeln!(
                s,
                "    c.v[{i}] = (a.v[{i}] - ({})) * inv_2c0;",
                terms.join(" + ")
            )
            .unwrap();
        }
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Sin/Cos coupled recurrence:
    // s[0] = sin(a[0]), co[0] = cos(a[0])
    // s[i] = (1/i) * Σ_{j=1}^{i} j*a[j]*co[i-j]
    // co[i] = -(1/i) * Σ_{j=1}^{i} j*a[j]*s[i-j]
    writeln!(s, "struct JetPair {{ a: JetK, b: JetK, }}\n").unwrap();

    writeln!(s, "fn jet_sin_cos(a: JetK) -> JetPair {{").unwrap();
    writeln!(s, "    var sn: JetK;").unwrap();
    writeln!(s, "    var co: JetK;").unwrap();
    writeln!(s, "    sn.v[0] = sin(a.v[0]);").unwrap();
    writeln!(s, "    co.v[0] = cos(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut sterms = Vec::new();
        let mut cterms = Vec::new();
        for j in 1..=i {
            sterms.push(format!("{:.1} * a.v[{j}] * co.v[{}]", j as f64, i - j));
            cterms.push(format!("{:.1} * a.v[{j}] * sn.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    sn.v[{i}] = {inv_i:.10} * ({});", sterms.join(" + ")).unwrap();
        writeln!(
            s,
            "    co.v[{i}] = -{inv_i:.10} * ({});",
            cterms.join(" + ")
        )
        .unwrap();
    }
    writeln!(s, "    return JetPair(sn, co);\n}}\n").unwrap();

    // Sinh/Cosh coupled: same but positive signs for cosh
    writeln!(s, "fn jet_sinh_cosh(a: JetK) -> JetPair {{").unwrap();
    writeln!(s, "    var sh: JetK;").unwrap();
    writeln!(s, "    var ch: JetK;").unwrap();
    writeln!(s, "    sh.v[0] = sinh_f(a.v[0]);").unwrap();
    writeln!(s, "    ch.v[0] = cosh_f(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut shterms = Vec::new();
        let mut chterms = Vec::new();
        for j in 1..=i {
            shterms.push(format!("{:.1} * a.v[{j}] * ch.v[{}]", j as f64, i - j));
            chterms.push(format!("{:.1} * a.v[{j}] * sh.v[{}]", j as f64, i - j));
        }
        writeln!(
            s,
            "    sh.v[{i}] = {inv_i:.10} * ({});",
            shterms.join(" + ")
        )
        .unwrap();
        writeln!(
            s,
            "    ch.v[{i}] = {inv_i:.10} * ({});",
            chterms.join(" + ")
        )
        .unwrap();
    }
    writeln!(s, "    return JetPair(sh, ch);\n}}\n").unwrap();

    // Tan: c' = a' * (1 + c²), where s = 1 + c² is a helper jet
    // s[0] = 1 + c[0]², then integration: c[i] = (1/i) * Σ j*a[j]*s[i-j]
    // But s depends on c, so we update s as we compute c.
    // s = jet_mul(c,c); s.v[0] += 1
    // c[i] = (1/i) * Σ_{j=1}^{i} j*a[j]*s[i-j], then update s[i] = Σ c[j]*c[i-j] for j=0..i
    writeln!(s, "fn jet_tan(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    var sc: JetK; // 1 + c²").unwrap();
    writeln!(s, "    c.v[0] = tan(a.v[0]);").unwrap();
    writeln!(s, "    sc.v[0] = 1.0 + c.v[0] * c.v[0];").unwrap();
    for i in 1..k {
        // First update sc[i] (c² contribution) from already-computed c[0..i-1]
        // Actually, we need sc[i-j] in the integration, so we compute sc up to i-1 first
        // Then compute c[i], then sc[i].
        // sc[p] = Σ_{j=0}^{p} c[j]*c[p-j] for the c² part
        // Actually the recurrence is: c[i] = (1/i) * Σ_{j=1}^{i} j*a[j]*sc[i-j]
        // where sc = 1+c². So sc[0] = 1 + c[0]², sc[p] = Σ_{j=0}^{p} c[j]*c[p-j] for p>0
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * sc.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = {inv_i:.10} * ({});", terms.join(" + ")).unwrap();
        // Now update sc[i] = Σ c[j]*c[i-j]
        let mut sc_terms = Vec::new();
        for j in 0..=i {
            sc_terms.push(format!("c.v[{j}] * c.v[{}]", i - j));
        }
        writeln!(s, "    sc.v[{i}] = {};", sc_terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Tanh: c' = a' * (1 - c²)
    writeln!(s, "fn jet_tanh(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    var sc: JetK; // 1 - c²").unwrap();
    writeln!(s, "    c.v[0] = tanh(a.v[0]);").unwrap();
    writeln!(s, "    sc.v[0] = 1.0 - c.v[0] * c.v[0];").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * sc.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = {inv_i:.10} * ({});", terms.join(" + ")).unwrap();
        let mut sc_terms = Vec::new();
        for j in 0..=i {
            sc_terms.push(format!("c.v[{j}] * c.v[{}]", i - j));
        }
        writeln!(s, "    sc.v[{i}] = -({});", sc_terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();
}

fn write_wgsl_jet_inverse_trig(s: &mut String, k: usize) {
    // atan: c' = a' / (1 + a²), so g = 1/(1+a²)
    writeln!(s, "fn jet_atan(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    let asq = jet_mul(a, a);").unwrap();
    write!(s, "    var d: JetK;\n    d.v[0] = 1.0 + asq.v[0];\n").unwrap();
    for i in 1..k {
        writeln!(s, "    d.v[{i}] = asq.v[{i}];").unwrap();
    }
    writeln!(s, "    let g = jet_recip(d);").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = atan(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * g.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = {inv_i:.10} * ({});", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // asin: c' = a' / sqrt(1 - a²)
    writeln!(s, "fn jet_asin(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    let asq = jet_mul(a, a);").unwrap();
    // Factored `(1-a0)(1+a0)` for the leading denom coeff avoids catastrophic
    // cancellation near |a0|=1 (matches CPU taylor_asin and jet_acosh); the
    // higher coeffs stay `-asq.v[i]` (bit-identical to the factored expansion).
    write!(
        s,
        "    var d: JetK;\n    d.v[0] = (1.0 - a.v[0]) * (1.0 + a.v[0]);\n"
    )
    .unwrap();
    for i in 1..k {
        writeln!(s, "    d.v[{i}] = -asq.v[{i}];").unwrap();
    }
    writeln!(s, "    let g = jet_recip(jet_sqrt(d));").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = asin(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * g.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = {inv_i:.10} * ({});", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // acos: π/2 - asin → negate higher coefficients
    writeln!(s, "fn jet_acos(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    let asq = jet_mul(a, a);").unwrap();
    // Factored leading denom coeff (same as jet_asin — acos shares the denom).
    write!(
        s,
        "    var d: JetK;\n    d.v[0] = (1.0 - a.v[0]) * (1.0 + a.v[0]);\n"
    )
    .unwrap();
    for i in 1..k {
        writeln!(s, "    d.v[{i}] = -asq.v[{i}];").unwrap();
    }
    writeln!(s, "    let g = jet_recip(jet_sqrt(d));").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = acos(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * g.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = -{inv_i:.10} * ({});", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // asinh: c' = a' / sqrt(1 + a²)
    writeln!(s, "fn jet_asinh(a: JetK) -> JetK {{").unwrap();
    writeln!(s, "    let asq = jet_mul(a, a);").unwrap();
    write!(s, "    var d: JetK;\n    d.v[0] = 1.0 + asq.v[0];\n").unwrap();
    for i in 1..k {
        writeln!(s, "    d.v[{i}] = asq.v[{i}];").unwrap();
    }
    writeln!(s, "    let g = jet_recip(jet_sqrt(d));").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = asinh_f(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * g.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = {inv_i:.10} * ({});", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // acosh: c' = a' / sqrt(a² - 1)
    // v[0] uses the factored form (a-1)·(a+1) to avoid catastrophic
    // cancellation near a=1 (matches kernels::acosh_deriv). v[i] for
    // i ≥ 1 stays as asq.v[i] — the higher-order Taylor coefficients
    // of (a-1)·(a+1) and a²-1 are mathematically identical (both
    // equal 2·a.v[0]·a.v[i] + Σ_{j+k=i, j,k≥1} a.v[j]·a.v[k]).
    writeln!(s, "fn jet_acosh(a: JetK) -> JetK {{").unwrap();
    // Domain guard: acosh(a0<1) is NaN. For a0<=-1 the factored (a0-1)(a0+1)>0
    // leaves a finite sqrt/recip → finite higher coeffs beside a NaN primal, so
    // guard explicitly (mirrors CPU taylor_acosh). Strict `< 1` keeps a0==1 as the
    // +Inf branch-point derivative.
    write_wgsl_jet_nan_guard(s, k, "a.v[0] < 1.0");
    writeln!(s, "    let asq = jet_mul(a, a);").unwrap();
    write!(
        s,
        "    var d: JetK;\n    d.v[0] = (a.v[0] - 1.0) * (a.v[0] + 1.0);\n"
    )
    .unwrap();
    for i in 1..k {
        writeln!(s, "    d.v[{i}] = asq.v[{i}];").unwrap();
    }
    writeln!(s, "    let g = jet_recip(jet_sqrt(d));").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = acosh_f(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * g.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = {inv_i:.10} * ({});", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // atanh: c' = a' / (1 - a²)
    writeln!(s, "fn jet_atanh(a: JetK) -> JetK {{").unwrap();
    // Domain guard: atanh(|a0|>1) is NaN, but recip(1-a0^2) stays finite there →
    // finite higher coeffs beside a NaN primal, so guard explicitly (mirrors CPU
    // taylor_atanh). Strict keeps |a0|==1 as the ±Inf branch-point derivative.
    write_wgsl_jet_nan_guard(s, k, "a.v[0] < -1.0 || a.v[0] > 1.0");
    writeln!(s, "    let asq = jet_mul(a, a);").unwrap();
    // Factored leading denom coeff (same as jet_asin — avoids cancellation near |a0|=1).
    write!(
        s,
        "    var d: JetK;\n    d.v[0] = (1.0 - a.v[0]) * (1.0 + a.v[0]);\n"
    )
    .unwrap();
    for i in 1..k {
        writeln!(s, "    d.v[{i}] = -asq.v[{i}];").unwrap();
    }
    writeln!(s, "    let g = jet_recip(d);").unwrap();
    writeln!(s, "    var c: JetK;").unwrap();
    writeln!(s, "    c.v[0] = atanh_f(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("{:.1} * a.v[{j}] * g.v[{}]", j as f64, i - j));
        }
        writeln!(s, "    c.v[{i}] = {inv_i:.10} * ({});", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();
}

/// Emit the WGSL Taylor-jet integer-power computation, mirroring CPU
/// `taylor_powi`: sets `r = a^exponent`, handling zero and negative bases
/// (which WGSL `pow` cannot). `n_expr` is an f32 exponent expression and
/// `ni_expr` its i32 form. Shared by the POWI opcode and the constant-
/// integer-exponent POWF path so the two cannot drift.
fn write_wgsl_powi_jet(s: &mut String, k: usize, n_expr: &str, ni_expr: &str) {
    writeln!(s, "                let n = {n_expr};").unwrap();
    writeln!(s, "                let ni = {ni_expr};").unwrap();
    writeln!(s, "                if n == 0.0 {{ r = jet_const(1.0); }}").unwrap();
    writeln!(s, "                else if n == 1.0 {{ r = a; }}").unwrap();
    // a.v[0] == 0 with small |n|: repeated squaring preserves the higher-order
    // Taylor coefficients (the log|a| = log(0) = -inf path below poisons them
    // with NaN). Chain matches CPU `taylor_powi_squaring`.
    writeln!(
        s,
        "                else if a.v[0] == 0.0 && ni == 2 {{ r = jet_mul(a, a); }}"
    )
    .unwrap();
    writeln!(
        s,
        "                else if a.v[0] == 0.0 && ni == 3 {{ r = jet_mul(jet_mul(a, a), a); }}"
    )
    .unwrap();
    writeln!(s, "                else if a.v[0] == 0.0 && ni == 4 {{ let a2 = jet_mul(a, a); r = jet_mul(a2, a2); }}").unwrap();
    writeln!(s, "                else if a.v[0] == 0.0 && ni == 5 {{ let a2 = jet_mul(a, a); let a4 = jet_mul(a2, a2); r = jet_mul(a4, a); }}").unwrap();
    writeln!(s, "                else if a.v[0] == 0.0 && ni == 6 {{ let a2 = jet_mul(a, a); let a4 = jet_mul(a2, a2); r = jet_mul(a4, a2); }}").unwrap();
    writeln!(s, "                else if a.v[0] == 0.0 && ni == 7 {{ let a2 = jet_mul(a, a); let a4 = jet_mul(a2, a2); r = jet_mul(jet_mul(a4, a2), a); }}").unwrap();
    writeln!(s, "                else if a.v[0] == 0.0 && ni == 8 {{ let a2 = jet_mul(a, a); let a4 = jet_mul(a2, a2); r = jet_mul(a4, a4); }}").unwrap();
    // Negative base (a.v[0] < 0) or a.v[0] == 0 with |n| > 8: sign(a)^n * |a|^n.
    writeln!(s, "                else if a.v[0] <= 0.0 {{").unwrap();
    writeln!(
        s,
        "                    let sf = select(1.0, -1.0, ni % 2 != 0);"
    )
    .unwrap();
    writeln!(
        s,
        "                    let sg = select(sign(a.v[0]), 1.0, a.v[0] == 0.0);"
    )
    .unwrap();
    writeln!(s, "                    var abs_a: JetK;").unwrap();
    writeln!(s, "                    abs_a.v[0] = abs(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "                    abs_a.v[{i}] = sg * a.v[{i}];").unwrap();
    }
    writeln!(
        s,
        "                    let e = jet_exp(jet_scale(jet_ln(abs_a), n));"
    )
    .unwrap();
    for i in 0..k {
        writeln!(s, "                    r.v[{i}] = sf * e.v[{i}];").unwrap();
    }
    writeln!(s, "                    r.v[0] = powf_real(a.v[0], n);").unwrap();
    writeln!(s, "                }} else {{").unwrap();
    writeln!(
        s,
        "                    r = jet_exp(jet_scale(jet_ln(a), n));"
    )
    .unwrap();
    writeln!(s, "                    r.v[0] = pow(a.v[0], n);").unwrap();
    writeln!(s, "                }}").unwrap();
}

fn write_wgsl_main_kernel(s: &mut String, k: usize) {
    writeln!(
        s,
        "@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {{
    let bid = gid.x;
    if bid >= tape_meta.batch_size {{
        return;
    }}

    let nv = tape_meta.num_variables;
    let ni = tape_meta.num_inputs;
    let num_ops = tape_meta.num_ops;
    let n_out = tape_meta.num_outputs;
    let K = {k}u;

    let j_base = bid * nv * K;"
    )
    .unwrap();

    // Initialize all variables from constants
    writeln!(s, "\n    // Initialize from constants").unwrap();
    writeln!(s, "    for (var i = 0u; i < nv; i = i + 1u) {{").unwrap();
    writeln!(s, "        let off = j_base + i * K;").unwrap();
    writeln!(s, "        jets[off] = constants[i];").unwrap();
    for c in 1..k {
        writeln!(s, "        jets[off + {c}u] = 0.0;").unwrap();
    }
    writeln!(s, "    }}").unwrap();

    // Set input variables
    writeln!(s, "\n    // Set input jets").unwrap();
    writeln!(s, "    let in_base = bid * ni;").unwrap();
    writeln!(s, "    for (var i = 0u; i < ni; i = i + 1u) {{").unwrap();
    writeln!(s, "        let off = j_base + i * K;").unwrap();
    writeln!(s, "        jets[off] = primal_inputs[in_base + i];").unwrap();
    if k > 1 {
        writeln!(s, "        jets[off + 1u] = direction_seeds[in_base + i];").unwrap();
    }
    writeln!(s, "    }}").unwrap();

    // Tape walk
    writeln!(s, "\n    // Walk the tape").unwrap();
    writeln!(s, "    for (var i = ni; i < num_ops; i = i + 1u) {{").unwrap();
    writeln!(s, "        let op = opcodes[i];").unwrap();
    writeln!(s, "        if op == OP_CONST {{ continue; }}").unwrap();
    writeln!(s, "        let a_idx = arg0[i];").unwrap();
    writeln!(s, "        let b_idx = arg1[i];").unwrap();
    writeln!(s, "        let a = jet_load(j_base + a_idx * K);").unwrap();
    writeln!(s, "        var r = jet_const(0.0);").unwrap();
    writeln!(s, "        switch op {{").unwrap();

    // Binary ops
    for (case, name) in &[
        (2, "jet_add"),
        (3, "jet_sub"),
        (4, "jet_mul"),
        (5, "jet_div"),
    ] {
        writeln!(s, "            case {case}u: {{").unwrap();
        writeln!(s, "                let b = jet_load(j_base + b_idx * K);").unwrap();
        writeln!(s, "                r = {name}(a, b);").unwrap();
        writeln!(s, "            }}").unwrap();
    }

    // REM: a % b = a - trunc(a[0]/b[0]) * b. Higher-order: r[k] = a[k] - q * b[k].
    writeln!(s, "            case 6u: {{").unwrap();
    writeln!(s, "                let b = jet_load(j_base + b_idx * K);").unwrap();
    // Zero-leading divisor → all-NaN jet (matches CPU Taylor::rem), not the
    // Inf/NaN mix that `trunc(x/0)` would produce.
    writeln!(s, "                if (b.v[0] == 0.0) {{").unwrap();
    for i in 0..k {
        writeln!(
            s,
            "                    r.v[{i}] = bitcast<f32>(0x7fc00000u);"
        )
        .unwrap();
    }
    writeln!(s, "                }} else {{").unwrap();
    writeln!(s, "                    let q = trunc(a.v[0] / b.v[0]);").unwrap();
    writeln!(s, "                    r.v[0] = a.v[0] - q * b.v[0];").unwrap();
    for i in 1..k {
        writeln!(s, "                    r.v[{i}] = a.v[{i}] - q * b.v[{i}];").unwrap();
    }
    writeln!(s, "                }}").unwrap();
    writeln!(s, "            }}").unwrap();

    // POWF — a constant integer exponent (in i32 range) is routed to the powi
    // jet, which is finite for any base (a<0 is well-defined at integer
    // powers), mirroring CPU `taylor_powf`. A non-integer or variable exponent
    // at a<=0 is genuinely NaN.
    writeln!(s, "            case 7u: {{").unwrap();
    writeln!(s, "                let b = jet_load(j_base + b_idx * K);").unwrap();
    // `< 2^31` (not `<= i32::MAX`): 2147483647.0 rounds up to 2^31 in f32, so
    // the `<=` form would admit 2^31 and overflow the `i32(..)` cast below.
    let mut b_int_cond = String::from("b.v[0] == round(b.v[0]) && abs(b.v[0]) < 2147483648.0");
    for i in 1..k {
        b_int_cond.push_str(&format!(" && b.v[{i}] == 0.0"));
    }
    writeln!(s, "                let b_is_int = {b_int_cond};").unwrap();
    writeln!(s, "                if b_is_int {{").unwrap();
    write_wgsl_powi_jet(s, k, "b.v[0]", "i32(b.v[0])");
    // Negative base with a non-constant-integer exponent: `a(t)^b(t)` is
    // complex for a varying exponent → all-NaN jet (matches CPU `taylor_powf`).
    // naga's `pow(neg, .)` is already NaN, but making it explicit keeps this in
    // lockstep with the CUDA emit, where C `pow(-2,3) = -8` would otherwise
    // leave a finite primal beside NaN derivative coefficients (a mixed jet).
    writeln!(s, "                }} else if a.v[0] < 0.0 {{").unwrap();
    for i in 0..k {
        writeln!(
            s,
            "                    r.v[{i}] = bitcast<f32>(0x7fc00000u);"
        )
        .unwrap();
    }
    writeln!(s, "                }} else if a.v[0] <= 0.0 {{").unwrap();
    writeln!(s, "                    let val = pow(a.v[0], b.v[0]);").unwrap();
    writeln!(
        s,
        "                    let da = b.v[0] * pow(a.v[0], b.v[0] - 1.0);"
    )
    .unwrap();
    // Reached only for a non-integer or variable exponent at a <= 0, where
    // a^b is genuinely NaN (real result undefined) — the constant-integer
    // case was routed to the powi jet above. db = 0 matches CPU OpCode::Powf.
    writeln!(s, "                    let db = 0.0;").unwrap();
    writeln!(s, "                    r.v[0] = val;").unwrap();
    if k > 1 {
        writeln!(s, "                    r.v[1] = da * a.v[1] + db * b.v[1];").unwrap();
    }
    for i in 2..k {
        // CPU powf with negative base produces NaN through exp(b*ln(negative)).
        // `0.0 / 0.0` is an abstract-typed NaN that naga rejects as not-
        // expressible in f32; use the canonical quiet-NaN bit pattern instead.
        writeln!(
            s,
            "                    r.v[{i}] = bitcast<f32>(0x7fc00000u);"
        )
        .unwrap();
    }
    writeln!(s, "                }} else {{").unwrap();
    writeln!(s, "                    let lna = jet_ln(a);").unwrap();
    writeln!(s, "                    let product = jet_mul(b, lna);").unwrap();
    writeln!(s, "                    r = jet_exp(product);").unwrap();
    writeln!(s, "                    r.v[0] = pow(a.v[0], b.v[0]);").unwrap();
    writeln!(s, "                }}").unwrap();
    writeln!(s, "            }}").unwrap();

    // ATAN2 — guard b==0 (division by zero in jet_div)
    writeln!(s, "            case 8u: {{").unwrap();
    writeln!(s, "                let b = jet_load(j_base + b_idx * K);").unwrap();
    writeln!(s, "                if b.v[0] == 0.0 {{").unwrap();
    writeln!(s, "                    r.v[0] = atan2(a.v[0], b.v[0]);").unwrap();
    writeln!(s, "                    if a.v[0] == 0.0 {{").unwrap();
    for i in 1..k {
        writeln!(s, "                        r.v[{i}] = 0.0;").unwrap();
    }
    writeln!(s, "                    }} else {{").unwrap();
    // b.v[0]==0, a.v[0]!=0: use atan2(a,b) = -atan(b/a) + pi/2*sign(a).
    // Compute full Taylor composition: ratio=b/a (ratio.v[0]=0), c=atan(ratio), negate.
    writeln!(s, "                        let ratio = jet_div(b, a);").unwrap();
    writeln!(s, "                        let c = jet_atan(ratio);").unwrap();
    for i in 1..k {
        writeln!(s, "                        r.v[{i}] = -c.v[{i}];").unwrap();
    }
    writeln!(s, "                    }}").unwrap();
    writeln!(s, "                }} else {{").unwrap();
    writeln!(s, "                    let ratio = jet_div(a, b);").unwrap();
    writeln!(s, "                    r = jet_atan(ratio);").unwrap();
    writeln!(s, "                    r.v[0] = atan2(a.v[0], b.v[0]);").unwrap();
    writeln!(s, "                }}").unwrap();
    writeln!(s, "            }}").unwrap();

    // HYPOT. Full jet-wide max-rescale to keep every coefficient
    // bounded even when `a.v[0] ~ 1e20`. Mirrors CPU
    // `taylor_ops::taylor_hypot`: build `a_s = a/h`, `b_s = b/h`,
    // compute `(a_s)² + (b_s)²`, sqrt, then scale back by `h`.
    // Branch order: Inf first, then NaN, then `h == 0` (with the
    // one-level shift-and-square unroll), then general — preserves
    // IEEE `hypot(Inf, NaN) = Inf` and `hypot(NaN, 0) = NaN`.
    //
    // Function-domain-boundary conventions (matching CPU):
    // - `hypot(Inf, finite)` — primal = Inf, higher-order = NaN.
    //   CPU's rescale path produces `Inf * 0 = NaN` in every
    //   coefficient past the primal; GPU now synthesises NaN via
    //   `inf - inf` (runtime) to match.
    // - `hypot(0, 0)` with `a.v[1] != 0 || b.v[1] != 0` —
    //   one-level shift-and-square unroll: compute `hypot(a/t, b/t)`
    //   via the rescale path on the shifted jet (at least one leading
    //   is now non-zero by construction), then shift result right by
    //   1 to represent the `|t|` factor in
    //   `t ↦ |t|·hypot(a/t, b/t)`. CPU recursion is at most one
    //   level deep since the entry check guarantees `scale > 0` in
    //   the recursive call.
    // - `hypot(0, 0)` with all higher-order seeds also zero —
    //   0 primal + Inf higher. Matches CPU's `taylor_sqrt` at a
    //   zero leading coefficient (produces Inf per
    //   `taylor_ops.rs:146-152`), then primal override gives 0.
    writeln!(s, "            case 9u: {{").unwrap();
    writeln!(s, "                let b = jet_load(j_base + b_idx * K);").unwrap();
    writeln!(s, "                let aa = abs(a.v[0]);").unwrap();
    writeln!(s, "                let bb = abs(b.v[0]);").unwrap();
    writeln!(s, "                let inf = bitcast<f32>(0x7f800000u);").unwrap();
    // Inf branch: `r` is initialised by the dispatcher (`var r =
    // jet_const(0.0)` at the top of the per-opcode loop) so the
    // explicit higher-order loop is *redundant* on WGSL — kept for
    // parity with CUDA (where `JetK r;` is uninitialised) and to
    // make the contract explicit. Higher-order emits NaN (runtime
    // via `inf - inf`) to match CPU's `Inf * 0 = NaN` rescale-path
    // output.
    writeln!(s, "                if (aa == inf || bb == inf) {{").unwrap();
    writeln!(s, "                    let nan = inf - inf;").unwrap();
    writeln!(s, "                    r.v[0] = inf;").unwrap();
    for i in 1..k {
        writeln!(s, "                    r.v[{i}] = nan;").unwrap();
    }
    // NaN detection: WGSL allows NaN comparisons to be "tame-folded"
    // (drivers may compile `x != x` to a constant false), so the
    // standard idiom is unreliable. Use a bit-pattern check on the
    // pre-`abs` values — for f32, NaN has exponent `0xff` and a
    // non-zero mantissa, so `(bits & 0x7fffffff) > 0x7f800000` is the
    // unambiguous test (Inf is exactly `0x7f800000`).
    writeln!(s, "                }} else if (").unwrap();
    writeln!(
        s,
        "                    (bitcast<u32>(a.v[0]) & 0x7fffffffu) > 0x7f800000u ||"
    )
    .unwrap();
    writeln!(
        s,
        "                    (bitcast<u32>(b.v[0]) & 0x7fffffffu) > 0x7f800000u"
    )
    .unwrap();
    writeln!(s, "                ) {{").unwrap();
    // NaN branch: without this, `max(NaN, 0) = 0` (Vulkan SPIR-V
    // `OpFMax` returns the non-NaN argument per IEEE 754-2008 maxNum),
    // causing a `hypot(NaN, 0)` input to fall into the `h == 0` branch
    // and return the zero jet — silently swallowing the NaN that IEEE
    // `hypot` requires propagating.
    writeln!(s, "                    r.v[0] = a.v[0] + b.v[0];").unwrap();
    for i in 1..k {
        writeln!(s, "                    r.v[{i}] = 0.0;").unwrap();
    }
    writeln!(s, "                }} else {{").unwrap();
    writeln!(s, "                    let h = max(aa, bb);").unwrap();
    // Zero branch (shift-and-square unroll).
    writeln!(s, "                    if (h == 0.0) {{").unwrap();
    if k >= 2 {
        // Build shifted jets: a_shifted.v[i] = a.v[i+1] for i < K-1,
        // a_shifted.v[K-1] = 0. Same for b. At least one of
        // a.v[1], b.v[1] is non-zero if we take the rescale branch;
        // else we emit 0 primal + Inf higher (matching CPU's
        // `taylor_sqrt` at a zero leading).
        writeln!(s, "                        var a_shifted: JetK;").unwrap();
        writeln!(s, "                        var b_shifted: JetK;").unwrap();
        for i in 0..(k - 1) {
            writeln!(
                s,
                "                        a_shifted.v[{i}] = a.v[{ip1}];",
                ip1 = i + 1
            )
            .unwrap();
            writeln!(
                s,
                "                        b_shifted.v[{i}] = b.v[{ip1}];",
                ip1 = i + 1
            )
            .unwrap();
        }
        writeln!(
            s,
            "                        a_shifted.v[{last}] = 0.0;",
            last = k - 1
        )
        .unwrap();
        writeln!(
            s,
            "                        b_shifted.v[{last}] = 0.0;",
            last = k - 1
        )
        .unwrap();
        writeln!(
            s,
            "                        let h_inner = max(abs(a_shifted.v[0]), abs(b_shifted.v[0]));"
        )
        .unwrap();
        writeln!(s, "                        if (h_inner == 0.0) {{").unwrap();
        // Deeper-order-zero: 0 primal + Inf higher.
        writeln!(s, "                            r.v[0] = 0.0;").unwrap();
        for i in 1..k {
            writeln!(s, "                            r.v[{i}] = inf;").unwrap();
        }
        writeln!(s, "                        }} else {{").unwrap();
        // Rescale + sum-of-squares + sqrt on the shifted jet.
        writeln!(
            s,
            "                            let inv_h_inner = 1.0 / h_inner;"
        )
        .unwrap();
        writeln!(
            s,
            "                            let a_ss = jet_scale(a_shifted, inv_h_inner);"
        )
        .unwrap();
        writeln!(
            s,
            "                            let b_ss = jet_scale(b_shifted, inv_h_inner);"
        )
        .unwrap();
        writeln!(
            s,
            "                            let sum_sq_s = jet_add(jet_mul(a_ss, a_ss), jet_mul(b_ss, b_ss));"
        )
        .unwrap();
        writeln!(
            s,
            "                            let r_s_inner = jet_sqrt(sum_sq_s);"
        )
        .unwrap();
        writeln!(
            s,
            "                            let inner = jet_scale(r_s_inner, h_inner);"
        )
        .unwrap();
        // Shift result right by 1: r.v[0] = 0, r.v[i] = inner.v[i-1].
        writeln!(s, "                            r.v[0] = 0.0;").unwrap();
        for i in 1..k {
            writeln!(
                s,
                "                            r.v[{i}] = inner.v[{im1}];",
                im1 = i - 1
            )
            .unwrap();
        }
        writeln!(s, "                        }}").unwrap();
    } else {
        // K == 1: no higher-order slots. Primal-only zero result.
        writeln!(s, "                        r.v[0] = 0.0;").unwrap();
    }
    writeln!(s, "                    }} else {{").unwrap();
    // General branch: rescale, sum-of-squares, sqrt, scale back.
    writeln!(s, "                        let inv_h = 1.0 / h;").unwrap();
    writeln!(s, "                        let a_s = jet_scale(a, inv_h);").unwrap();
    writeln!(s, "                        let b_s = jet_scale(b, inv_h);").unwrap();
    writeln!(
        s,
        "                        let sum_sq = jet_add(jet_mul(a_s, a_s), jet_mul(b_s, b_s));"
    )
    .unwrap();
    writeln!(s, "                        let r_s = jet_sqrt(sum_sq);").unwrap();
    writeln!(s, "                        r = jet_scale(r_s, h);").unwrap();
    // Primal override: `r.v[0] = jet_scale(r_s, h).v[0]` is bit-
    // equivalent to the explicit rescale formula below (same arithmetic,
    // same rounding under default WGSL contraction settings), so this
    // assignment is technically redundant. Keep it as a regression
    // guard against future jet-helper refactors that might change
    // rounding behaviour.
    writeln!(
        s,
        "                        let a_s0 = a.v[0] * inv_h; let b_s0 = b.v[0] * inv_h;"
    )
    .unwrap();
    writeln!(
        s,
        "                        r.v[0] = h * sqrt(a_s0 * a_s0 + b_s0 * b_s0);"
    )
    .unwrap();
    writeln!(s, "                    }}").unwrap();
    writeln!(s, "                }}").unwrap();
    writeln!(s, "            }}").unwrap();

    // MAX, MIN
    writeln!(s, "            case 10u: {{").unwrap();
    writeln!(s, "                let b = jet_load(j_base + b_idx * K);").unwrap();
    writeln!(
        s,
        "                if a.v[0] >= b.v[0] || is_nan_f32(b.v[0]) {{ r = a; }} else {{ r = b; }}"
    )
    .unwrap();
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            case 11u: {{").unwrap();
    writeln!(s, "                let b = jet_load(j_base + b_idx * K);").unwrap();
    writeln!(
        s,
        "                if a.v[0] <= b.v[0] || is_nan_f32(b.v[0]) {{ r = a; }} else {{ r = b; }}"
    )
    .unwrap();
    writeln!(s, "            }}").unwrap();

    // Unary ops
    writeln!(s, "            case 12u: {{ r = jet_neg(a); }}").unwrap();
    writeln!(s, "            case 13u: {{ r = jet_recip(a); }}").unwrap();
    writeln!(s, "            case 14u: {{ r = jet_sqrt(a); }}").unwrap();

    // CBRT — guard a.v[0]==0 (vertical tangent singularity)
    writeln!(s, "            case 15u: {{").unwrap();
    writeln!(s, "                if a.v[0] == 0.0 {{").unwrap();
    writeln!(s, "                    r.v[0] = 0.0;").unwrap();
    for i in 1..k {
        // `1.0 / 0.0` is a compile-time abstract `inf` that naga rejects as
        // inexpressible in f32. Use the bit pattern for +∞ instead.
        writeln!(
            s,
            "                    r.v[{i}] = bitcast<f32>(0x7f800000u);"
        )
        .unwrap();
    }
    writeln!(s, "                }} else {{").unwrap();
    writeln!(
        s,
        "                    let sg = select(sign(a.v[0]), 1.0, a.v[0] == 0.0);"
    )
    .unwrap();
    write!(
        s,
        "                    var abs_a: JetK;\n                    abs_a.v[0] = abs(a.v[0]);\n"
    )
    .unwrap();
    for i in 1..k {
        writeln!(s, "                    abs_a.v[{i}] = sg * a.v[{i}];").unwrap();
    }
    writeln!(s, "                    let lna = jet_ln(abs_a);").unwrap();
    writeln!(
        s,
        "                    let third = jet_scale(lna, 1.0 / 3.0);"
    )
    .unwrap();
    writeln!(s, "                    let e = jet_exp(third);").unwrap();
    writeln!(s, "                    r.v[0] = sg * e.v[0];").unwrap();
    for i in 1..k {
        writeln!(s, "                    r.v[{i}] = sg * e.v[{i}];").unwrap();
    }
    writeln!(s, "                }}").unwrap();
    writeln!(s, "            }}").unwrap();

    // POWI — handle negative bases via sign(a)^n * exp(n * ln(|a|))
    writeln!(s, "            case 16u: {{").unwrap();
    write_wgsl_powi_jet(s, k, "f32(bitcast<i32>(b_idx))", "bitcast<i32>(b_idx)");
    writeln!(s, "            }}").unwrap();

    // Transcendental unary
    writeln!(s, "            case 17u: {{ r = jet_exp(a); }}").unwrap();
    writeln!(s, "            case 18u: {{").unwrap();
    writeln!(s, "                let ln2 = log(2.0);").unwrap();
    writeln!(s, "                let scaled = jet_scale(a, ln2);").unwrap();
    writeln!(s, "                r = jet_exp(scaled);").unwrap();
    writeln!(s, "                r.v[0] = exp2(a.v[0]);").unwrap();
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            case 19u: {{").unwrap();
    writeln!(s, "                r = jet_exp(a);").unwrap();
    writeln!(s, "                r.v[0] = expm1_f(a.v[0]);").unwrap();
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            case 20u: {{ r = jet_ln(a); }}").unwrap();
    writeln!(s, "            case 21u: {{").unwrap();
    writeln!(s, "                r = jet_ln(a);").unwrap();
    writeln!(s, "                let inv_ln2 = 1.0 / log(2.0);").unwrap();
    writeln!(s, "                r.v[0] = log2(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "                r.v[{i}] = r.v[{i}] * inv_ln2;").unwrap();
    }
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            case 22u: {{").unwrap();
    writeln!(s, "                r = jet_ln(a);").unwrap();
    writeln!(s, "                let inv_ln10 = 1.0 / log(10.0);").unwrap();
    writeln!(s, "                r.v[0] = log(a.v[0]) * inv_ln10;").unwrap();
    for i in 1..k {
        writeln!(s, "                r.v[{i}] = r.v[{i}] * inv_ln10;").unwrap();
    }
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            case 23u: {{").unwrap();
    write!(
        s,
        "                var one_plus_a: JetK;\n                one_plus_a.v[0] = 1.0 + a.v[0];\n"
    )
    .unwrap();
    for i in 1..k {
        writeln!(s, "                one_plus_a.v[{i}] = a.v[{i}];").unwrap();
    }
    writeln!(s, "                r = jet_ln(one_plus_a);").unwrap();
    writeln!(s, "                r.v[0] = log1p_f(a.v[0]);").unwrap();
    writeln!(s, "            }}").unwrap();

    // Sin, Cos
    writeln!(
        s,
        "            case 24u: {{ let sc = jet_sin_cos(a); r = sc.a; }}"
    )
    .unwrap();
    writeln!(
        s,
        "            case 25u: {{ let sc = jet_sin_cos(a); r = sc.b; }}"
    )
    .unwrap();
    writeln!(s, "            case 26u: {{ r = jet_tan(a); }}").unwrap();
    writeln!(s, "            case 27u: {{ r = jet_asin(a); }}").unwrap();
    writeln!(s, "            case 28u: {{ r = jet_acos(a); }}").unwrap();
    writeln!(s, "            case 29u: {{ r = jet_atan(a); }}").unwrap();
    writeln!(
        s,
        "            case 30u: {{ let sc = jet_sinh_cosh(a); r = sc.a; }}"
    )
    .unwrap();
    writeln!(
        s,
        "            case 31u: {{ let sc = jet_sinh_cosh(a); r = sc.b; }}"
    )
    .unwrap();
    writeln!(s, "            case 32u: {{ r = jet_tanh(a); }}").unwrap();
    writeln!(s, "            case 33u: {{ r = jet_asinh(a); }}").unwrap();
    writeln!(s, "            case 34u: {{ r = jet_acosh(a); }}").unwrap();
    writeln!(s, "            case 35u: {{ r = jet_atanh(a); }}").unwrap();

    // Nonsmooth — ABS: scan for the first non-zero Taylor coefficient to
    // decide the sign. Matches CPU `Taylor::abs` semantics: at a = [0, -v, …]
    // the function |f(t)| ≈ v·|t| in a neighborhood of t=0, so the effective
    // sign comes from the first-non-zero coefficient (here -v → sg = -1).
    // Defaulting sg = +1 at a.v[0]=0 (as `sign(0) = 0` would prompt) gives
    // the wrong branch when a.v[1] < 0 or later coefficients dominate.
    writeln!(s, "            case 36u: {{").unwrap();
    writeln!(s, "                var sg: f32 = 1.0;").unwrap();
    for i in 0..k {
        writeln!(
            s,
            "                if a.v[{i}] != 0.0 {{ sg = select(1.0, -1.0, a.v[{i}] < 0.0); }} else",
        )
        .unwrap();
    }
    // Close the else-if chain with an empty else block; sg stays at 1.0 for
    // the identically-zero series (Rust convention).
    writeln!(s, "                {{ }}").unwrap();
    writeln!(s, "                r.v[0] = abs(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "                r.v[{i}] = sg * a.v[{i}];").unwrap();
    }
    writeln!(s, "            }}").unwrap();

    writeln!(s, "            case 37u, 38u, 39u, 40u, 41u: {{").unwrap();
    writeln!(s, "                var val = 0.0f;").unwrap();
    writeln!(s, "                switch op {{").unwrap();
    writeln!(
        s,
        "                    case 37u: {{ val = signum_f32(a.v[0]); }}"
    )
    .unwrap();
    writeln!(
        s,
        "                    case 38u: {{ val = floor(a.v[0]); }}"
    )
    .unwrap();
    writeln!(s, "                    case 39u: {{ val = ceil(a.v[0]); }}").unwrap();
    writeln!(
        s,
        "                    case 40u: {{ let t = trunc(a.v[0]); val = select(t, t + select(-1.0, 1.0, a.v[0] >= 0.0), abs(a.v[0] - t) >= 0.5); }}"
    )
    .unwrap();
    writeln!(
        s,
        "                    case 41u: {{ val = trunc(a.v[0]); }}"
    )
    .unwrap();
    writeln!(s, "                    default: {{}}").unwrap();
    writeln!(s, "                }}").unwrap();
    writeln!(s, "                r = jet_const(val);").unwrap();
    writeln!(s, "            }}").unwrap();

    writeln!(s, "            case 42u: {{").unwrap();
    writeln!(s, "                r.v[0] = fract(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "                r.v[{i}] = a.v[{i}];").unwrap();
    }
    writeln!(s, "            }}").unwrap();

    writeln!(s, "            default: {{}}").unwrap();
    writeln!(s, "        }}").unwrap();

    // Store result
    writeln!(s, "        jet_store(j_base + i * K, r);").unwrap();
    writeln!(s, "    }}").unwrap();

    // Write output jets
    writeln!(s, "\n    // Write output jets").unwrap();
    writeln!(s, "    let out_base = bid * n_out * K;").unwrap();
    writeln!(s, "    for (var j = 0u; j < n_out; j = j + 1u) {{").unwrap();
    writeln!(s, "        let oi = output_indices[j];").unwrap();
    writeln!(s, "        let src = j_base + oi * K;").unwrap();
    writeln!(s, "        let dst = out_base + j * K;").unwrap();
    for c in 0..k {
        writeln!(s, "        jet_outputs[dst + {c}u] = jets[src + {c}u];").unwrap();
    }
    writeln!(s, "    }}").unwrap();
    writeln!(s, "}}").unwrap();
}

// ══════════════════════════════════════════════
//  CUDA code generation
// ══════════════════════════════════════════════

fn write_cuda_opcodes(s: &mut String) {
    let ops = [
        "OP_INPUT",
        "OP_CONST",
        "OP_ADD",
        "OP_SUB",
        "OP_MUL",
        "OP_DIV",
        "OP_REM",
        "OP_POWF",
        "OP_ATAN2",
        "OP_HYPOT",
        "OP_MAX",
        "OP_MIN",
        "OP_NEG",
        "OP_RECIP",
        "OP_SQRT",
        "OP_CBRT",
        "OP_POWI",
        "OP_EXP",
        "OP_EXP2",
        "OP_EXPM1",
        "OP_LN",
        "OP_LOG2",
        "OP_LOG10",
        "OP_LN1P",
        "OP_SIN",
        "OP_COS",
        "OP_TAN",
        "OP_ASIN",
        "OP_ACOS",
        "OP_ATAN",
        "OP_SINH",
        "OP_COSH",
        "OP_TANH",
        "OP_ASINH",
        "OP_ACOSH",
        "OP_ATANH",
        "OP_ABS",
        "OP_SIGNUM",
        "OP_FLOOR",
        "OP_CEIL",
        "OP_ROUND",
        "OP_TRUNC",
        "OP_FRACT",
    ];
    for (i, name) in ops.iter().enumerate() {
        writeln!(s, "#define {name} {i}").unwrap();
    }
    writeln!(s).unwrap();
}

fn write_cuda_helpers(s: &mut String) {
    // _sign matches Rust's f64::signum: +1 for +0, -1 for -0, NaN for NaN.
    writeln!(
        s,
        "__device__ F _sign(F x) {{ return (x != x) ? x : copysign((F)1, x); }}
__device__ F _cbrt_f(F x) {{ return (x >= (F)0) ? pow(x, (F)(1.0/3.0)) : -pow(-x, (F)(1.0/3.0)); }}
__device__ F _fract(F x) {{ return x - floor(x); }}
"
    )
    .unwrap();
}

fn write_cuda_jet_type(s: &mut String, k: usize) {
    writeln!(
        s,
        "struct JetK {{
    F v[{k}];
    __device__ JetK() {{ for(int i=0;i<{k};i++) v[i]=(F)0; }}
}};
"
    )
    .unwrap();

    // jet_const
    writeln!(s, "__device__ JetK jet_const(F val) {{").unwrap();
    writeln!(s, "    JetK j;").unwrap();
    writeln!(s, "    j.v[0] = val;").unwrap();
    writeln!(s, "    return j;\n}}\n").unwrap();
}

fn write_cuda_jet_arithmetic(s: &mut String, k: usize) {
    // Add
    writeln!(s, "__device__ JetK jet_add(JetK a, JetK b) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    for i in 0..k {
        writeln!(s, "    c.v[{i}] = a.v[{i}] + b.v[{i}];").unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Sub
    writeln!(s, "__device__ JetK jet_sub(JetK a, JetK b) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    for i in 0..k {
        writeln!(s, "    c.v[{i}] = a.v[{i}] - b.v[{i}];").unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Neg
    writeln!(s, "__device__ JetK jet_neg(JetK a) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    for i in 0..k {
        writeln!(s, "    c.v[{i}] = -a.v[{i}];").unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Scale
    writeln!(s, "__device__ JetK jet_scale(JetK a, F s) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    for i in 0..k {
        writeln!(s, "    c.v[{i}] = a.v[{i}] * s;").unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Mul (Cauchy product)
    writeln!(s, "__device__ JetK jet_mul(JetK a, JetK b) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    for i in 0..k {
        let mut terms = Vec::new();
        for j in 0..=i {
            terms.push(format!("a.v[{j}] * b.v[{}]", i - j));
        }
        writeln!(s, "    c.v[{i}] = {};", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Div
    writeln!(s, "__device__ JetK jet_div(JetK a, JetK b) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    writeln!(s, "    F inv_b0 = (F)1 / b.v[0];").unwrap();
    writeln!(s, "    c.v[0] = a.v[0] * inv_b0;").unwrap();
    for i in 1..k {
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("b.v[{j}] * c.v[{}]", i - j));
        }
        writeln!(
            s,
            "    c.v[{i}] = (a.v[{i}] - ({})) * inv_b0;",
            terms.join(" + ")
        )
        .unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Recip
    writeln!(s, "__device__ JetK jet_recip(JetK a) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    writeln!(s, "    c.v[0] = (F)1 / a.v[0];").unwrap();
    for i in 1..k {
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("a.v[{j}] * c.v[{}]", i - j));
        }
        writeln!(s, "    c.v[{i}] = -({}) * c.v[0];", terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();
}

fn write_cuda_jet_transcendental(s: &mut String, k: usize) {
    // Exp
    writeln!(s, "__device__ JetK jet_exp(JetK a) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    writeln!(s, "    c.v[0] = exp(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("(F){:.1} * a.v[{j}] * c.v[{}]", j as f64, i - j));
        }
        writeln!(
            s,
            "    c.v[{i}] = (F){inv_i:.10} * ({});",
            terms.join(" + ")
        )
        .unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Ln
    writeln!(s, "__device__ JetK jet_ln(JetK a) {{").unwrap();
    // Domain guard (mirrors CPU taylor_ln; covers log2/log10/ln1p by delegation).
    write_cuda_jet_nan_guard(s, k, "a.v[0] < (F)0");
    writeln!(s, "    JetK c;").unwrap();
    writeln!(s, "    c.v[0] = log(a.v[0]);").unwrap();
    // `inv_a0` and the recurrence only exist for K > 1; emitting the
    // declaration at K=1 (empty loop) leaves an unused variable in the kernel.
    if k > 1 {
        writeln!(s, "    F inv_a0 = (F)1 / a.v[0];").unwrap();
        for i in 1..k {
            if i == 1 {
                writeln!(s, "    c.v[1] = a.v[1] * inv_a0;").unwrap();
            } else {
                let inv_i = 1.0 / i as f64;
                let mut terms = Vec::new();
                for j in 1..i {
                    terms.push(format!("(F){:.1} * c.v[{j}] * a.v[{}]", j as f64, i - j));
                }
                writeln!(
                    s,
                    "    c.v[{i}] = (a.v[{i}] - (F){inv_i:.10} * ({})) * inv_a0;",
                    terms.join(" + ")
                )
                .unwrap();
            }
        }
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Sqrt
    writeln!(s, "__device__ JetK jet_sqrt(JetK a) {{").unwrap();
    writeln!(s, "    JetK c;").unwrap();
    // Uniform +Inf higher coeffs at a0==0 (vertical tangent), matching CPU taylor_sqrt.
    writeln!(s, "    if (a.v[0] == (F)0) {{").unwrap();
    writeln!(s, "        c.v[0] = (F)0;").unwrap();
    for i in 1..k {
        writeln!(s, "        c.v[{i}] = F(1.0/0.0);").unwrap();
    }
    writeln!(s, "        return c;").unwrap();
    writeln!(s, "    }}").unwrap();
    writeln!(s, "    c.v[0] = sqrt(a.v[0]);").unwrap();
    if k > 1 {
        writeln!(s, "    F inv_2c0 = (F)0.5 / c.v[0];").unwrap();
    }
    for i in 1..k {
        if i == 1 {
            writeln!(s, "    c.v[1] = a.v[1] * inv_2c0;").unwrap();
        } else {
            let mut terms = Vec::new();
            for j in 1..i {
                terms.push(format!("c.v[{j}] * c.v[{}]", i - j));
            }
            writeln!(
                s,
                "    c.v[{i}] = (a.v[{i}] - ({})) * inv_2c0;",
                terms.join(" + ")
            )
            .unwrap();
        }
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Sin/Cos
    writeln!(s, "struct JetPair {{ JetK a; JetK b; }};\n").unwrap();

    writeln!(s, "__device__ JetPair jet_sin_cos(JetK a) {{").unwrap();
    writeln!(s, "    JetK sn, co;").unwrap();
    writeln!(s, "    sn.v[0] = sin(a.v[0]);").unwrap();
    writeln!(s, "    co.v[0] = cos(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut sterms = Vec::new();
        let mut cterms = Vec::new();
        for j in 1..=i {
            sterms.push(format!("(F){:.1} * a.v[{j}] * co.v[{}]", j as f64, i - j));
            cterms.push(format!("(F){:.1} * a.v[{j}] * sn.v[{}]", j as f64, i - j));
        }
        writeln!(
            s,
            "    sn.v[{i}] = (F){inv_i:.10} * ({});",
            sterms.join(" + ")
        )
        .unwrap();
        writeln!(
            s,
            "    co.v[{i}] = -(F){inv_i:.10} * ({});",
            cterms.join(" + ")
        )
        .unwrap();
    }
    writeln!(s, "    JetPair p; p.a = sn; p.b = co;").unwrap();
    writeln!(s, "    return p;\n}}\n").unwrap();

    // Sinh/Cosh
    writeln!(s, "__device__ JetPair jet_sinh_cosh(JetK a) {{").unwrap();
    writeln!(s, "    JetK sh, ch;").unwrap();
    writeln!(s, "    sh.v[0] = sinh(a.v[0]);").unwrap();
    writeln!(s, "    ch.v[0] = cosh(a.v[0]);").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut shterms = Vec::new();
        let mut chterms = Vec::new();
        for j in 1..=i {
            shterms.push(format!("(F){:.1} * a.v[{j}] * ch.v[{}]", j as f64, i - j));
            chterms.push(format!("(F){:.1} * a.v[{j}] * sh.v[{}]", j as f64, i - j));
        }
        writeln!(
            s,
            "    sh.v[{i}] = (F){inv_i:.10} * ({});",
            shterms.join(" + ")
        )
        .unwrap();
        writeln!(
            s,
            "    ch.v[{i}] = (F){inv_i:.10} * ({});",
            chterms.join(" + ")
        )
        .unwrap();
    }
    writeln!(s, "    JetPair p; p.a = sh; p.b = ch;").unwrap();
    writeln!(s, "    return p;\n}}\n").unwrap();

    // Tan
    writeln!(s, "__device__ JetK jet_tan(JetK a) {{").unwrap();
    writeln!(s, "    JetK c, sc;").unwrap();
    writeln!(s, "    c.v[0] = tan(a.v[0]);").unwrap();
    writeln!(s, "    sc.v[0] = (F)1 + c.v[0] * c.v[0];").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("(F){:.1} * a.v[{j}] * sc.v[{}]", j as f64, i - j));
        }
        writeln!(
            s,
            "    c.v[{i}] = (F){inv_i:.10} * ({});",
            terms.join(" + ")
        )
        .unwrap();
        let mut sc_terms = Vec::new();
        for j in 0..=i {
            sc_terms.push(format!("c.v[{j}] * c.v[{}]", i - j));
        }
        writeln!(s, "    sc.v[{i}] = {};", sc_terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();

    // Tanh
    writeln!(s, "__device__ JetK jet_tanh(JetK a) {{").unwrap();
    writeln!(s, "    JetK c, sc;").unwrap();
    writeln!(s, "    c.v[0] = tanh(a.v[0]);").unwrap();
    writeln!(s, "    sc.v[0] = (F)1 - c.v[0] * c.v[0];").unwrap();
    for i in 1..k {
        let inv_i = 1.0 / i as f64;
        let mut terms = Vec::new();
        for j in 1..=i {
            terms.push(format!("(F){:.1} * a.v[{j}] * sc.v[{}]", j as f64, i - j));
        }
        writeln!(
            s,
            "    c.v[{i}] = (F){inv_i:.10} * ({});",
            terms.join(" + ")
        )
        .unwrap();
        let mut sc_terms = Vec::new();
        for j in 0..=i {
            sc_terms.push(format!("c.v[{j}] * c.v[{}]", i - j));
        }
        writeln!(s, "    sc.v[{i}] = -({});", sc_terms.join(" + ")).unwrap();
    }
    writeln!(s, "    return c;\n}}\n").unwrap();
}

fn write_cuda_jet_inverse_trig(s: &mut String, k: usize) {
    // Helper macro for inverse trig/hyp functions
    // All follow the pattern: compute derivative factor g, then integrate
    let inv_trig_fns = [
        ("atan", "1.0 + asq", false, "atan(a.v[0])", false), // g = 1/(1+a²)
        ("asin", "1.0 - asq", true, "asin(a.v[0])", false),  // g = 1/sqrt(1-a²)
        ("acos", "1.0 - asq", true, "acos(a.v[0])", true),   // g = -1/sqrt(1-a²)
        ("asinh", "1.0 + asq", true, "asinh(a.v[0])", false), // g = 1/sqrt(1+a²)
        ("acosh", "asq - 1.0", true, "acosh(a.v[0])", false), // g = 1/sqrt(a²-1)
        ("atanh", "1.0 - asq", false, "atanh(a.v[0])", false), // g = 1/(1-a²)
    ];

    for (name, d_expr, use_sqrt, c0_expr, negate) in &inv_trig_fns {
        writeln!(s, "__device__ JetK jet_{name}(JetK a) {{").unwrap();
        // Series domain-NaN guard for the domain-restricted inverse fns (mirrors CPU
        // taylor_acosh/taylor_atanh). asin/acos self-NaN via jet_sqrt(negative), and
        // atan/asinh/acosh(via factoring)... only acosh (a0<1) and atanh (|a0|>1)
        // leave finite higher coeffs out of domain, so only they need an explicit guard.
        match *name {
            "acosh" => write_cuda_jet_nan_guard(s, k, "a.v[0] < (F)1"),
            "atanh" => write_cuda_jet_nan_guard(s, k, "a.v[0] < (F)-1 || a.v[0] > (F)1"),
            _ => {}
        }
        writeln!(s, "    JetK asq = jet_mul(a, a);").unwrap();
        writeln!(s, "    JetK d;").unwrap();

        // Build d from expression template
        if d_expr.starts_with("1.0 +") {
            writeln!(s, "    d.v[0] = (F)1 + asq.v[0];").unwrap();
            for i in 1..k {
                writeln!(s, "    d.v[{i}] = asq.v[{i}];").unwrap();
            }
        } else if d_expr.starts_with("1.0 -") {
            // Factored `(1-a0)(1+a0)` for the leading denom coeff avoids
            // cancellation near |a0|=1 (matches CPU taylor_asin/taylor_atanh and
            // the acosh branch below); asin/acos/atanh all share this branch.
            writeln!(s, "    d.v[0] = ((F)1 - a.v[0]) * ((F)1 + a.v[0]);").unwrap();
            for i in 1..k {
                writeln!(s, "    d.v[{i}] = -asq.v[{i}];").unwrap();
            }
        } else {
            // asq - 1.0 — emit factored form for v[0] to retain precision
            // near a=1; matches kernels::acosh_deriv. Higher-order
            // coefficients are unchanged (mathematically identical between
            // the two forms — see jet_acosh emitter for the full argument).
            writeln!(s, "    d.v[0] = (a.v[0] - (F)1) * (a.v[0] + (F)1);").unwrap();
            for i in 1..k {
                writeln!(s, "    d.v[{i}] = asq.v[{i}];").unwrap();
            }
        }

        if *use_sqrt {
            writeln!(s, "    JetK g = jet_recip(jet_sqrt(d));").unwrap();
        } else {
            writeln!(s, "    JetK g = jet_recip(d);").unwrap();
        }

        writeln!(s, "    JetK c;").unwrap();
        writeln!(s, "    c.v[0] = {c0_expr};").unwrap();
        let sign_str = if *negate { "-" } else { "" };
        for i in 1..k {
            let inv_i = 1.0 / i as f64;
            let mut terms = Vec::new();
            for j in 1..=i {
                terms.push(format!("(F){:.1} * a.v[{j}] * g.v[{}]", j as f64, i - j));
            }
            writeln!(
                s,
                "    c.v[{i}] = {sign_str}(F){inv_i:.10} * ({});",
                terms.join(" + ")
            )
            .unwrap();
        }
        writeln!(s, "    return c;\n}}\n").unwrap();
    }
}

/// Emit the CUDA Taylor-jet integer-power computation, mirroring CPU
/// `taylor_powi` (and the WGSL `write_wgsl_powi_jet`): sets `r = a^exponent`,
/// handling zero and negative bases. `n_expr` is an F exponent expression and
/// `ni_expr` its int form. Shared by the POWI opcode and the constant-integer-
/// exponent POWF path so the two cannot drift.
fn write_cuda_powi_jet(s: &mut String, k: usize, n_expr: &str, ni_expr: &str) {
    writeln!(s, "            int ni = {ni_expr};").unwrap();
    writeln!(s, "            F n = {n_expr};").unwrap();
    writeln!(s, "            if (ni == 0) {{ r = jet_const((F)1); }}").unwrap();
    writeln!(s, "            else if (ni == 1) {{ r = a; }}").unwrap();
    // a.v[0] == 0 with small |n|: repeated squaring preserves higher-order
    // coefficients (the log|a| = log(0) = -inf path poisons them). Matches CPU
    // `taylor_powi_squaring`.
    writeln!(
        s,
        "            else if (a.v[0] == F(0) && ni == 2) {{ r = jet_mul(a, a); }}"
    )
    .unwrap();
    writeln!(
        s,
        "            else if (a.v[0] == F(0) && ni == 3) {{ r = jet_mul(jet_mul(a, a), a); }}"
    )
    .unwrap();
    writeln!(s, "            else if (a.v[0] == F(0) && ni == 4) {{ JetK a2 = jet_mul(a, a); r = jet_mul(a2, a2); }}").unwrap();
    writeln!(s, "            else if (a.v[0] == F(0) && ni == 5) {{ JetK a2 = jet_mul(a, a); JetK a4 = jet_mul(a2, a2); r = jet_mul(a4, a); }}").unwrap();
    writeln!(s, "            else if (a.v[0] == F(0) && ni == 6) {{ JetK a2 = jet_mul(a, a); JetK a4 = jet_mul(a2, a2); r = jet_mul(a4, a2); }}").unwrap();
    writeln!(s, "            else if (a.v[0] == F(0) && ni == 7) {{ JetK a2 = jet_mul(a, a); JetK a4 = jet_mul(a2, a2); r = jet_mul(jet_mul(a4, a2), a); }}").unwrap();
    writeln!(s, "            else if (a.v[0] == F(0) && ni == 8) {{ JetK a2 = jet_mul(a, a); JetK a4 = jet_mul(a2, a2); r = jet_mul(a4, a4); }}").unwrap();
    // Negative base (or a==0 with |n|>8): sign(a)^n * |a|^n. C `pow` handles a
    // negative base at an integer exponent (Annex F), so r.v[0] stays native.
    writeln!(s, "            else if (a.v[0] <= F(0)) {{").unwrap();
    writeln!(s, "                F sf = (ni % 2 == 0) ? F(1) : F(-1);").unwrap();
    writeln!(s, "                F sg = _sign(a.v[0]);").unwrap();
    writeln!(s, "                JetK abs_a; abs_a.v[0] = fabs(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "                abs_a.v[{i}] = sg * a.v[{i}];").unwrap();
    }
    writeln!(
        s,
        "                JetK e = jet_exp(jet_scale(jet_ln(abs_a), n));"
    )
    .unwrap();
    for i in 0..k {
        writeln!(s, "                r.v[{i}] = sf * e.v[{i}];").unwrap();
    }
    writeln!(s, "                r.v[0] = pow(a.v[0], n);").unwrap();
    writeln!(s, "            }}").unwrap();
    writeln!(
        s,
        "            else {{ r = jet_exp(jet_scale(jet_ln(a), n)); r.v[0] = pow(a.v[0], n); }}"
    )
    .unwrap();
}

fn write_cuda_main_kernel(s: &mut String, k: usize) {
    writeln!(s, "extern \"C\" __global__ void taylor_forward_kth(").unwrap();
    writeln!(s, "    const unsigned int* __restrict__ opcodes,").unwrap();
    writeln!(s, "    const unsigned int* __restrict__ arg0,").unwrap();
    writeln!(s, "    const unsigned int* __restrict__ arg1,").unwrap();
    writeln!(s, "    const F* __restrict__ constants,").unwrap();
    writeln!(s, "    const F* __restrict__ primal_inputs,").unwrap();
    writeln!(s, "    const F* __restrict__ direction_seeds,").unwrap();
    writeln!(s, "    F* __restrict__ jets,").unwrap();
    writeln!(s, "    F* __restrict__ jet_outputs,").unwrap();
    writeln!(s, "    const unsigned int* __restrict__ output_indices,").unwrap();
    writeln!(s, "    unsigned int num_ops,").unwrap();
    writeln!(s, "    unsigned int num_inputs,").unwrap();
    writeln!(s, "    unsigned int num_variables,").unwrap();
    writeln!(s, "    unsigned int num_outputs,").unwrap();
    writeln!(s, "    unsigned int batch_size").unwrap();
    writeln!(s, ") {{").unwrap();
    writeln!(
        s,
        "    unsigned int bid = blockIdx.x * blockDim.x + threadIdx.x;"
    )
    .unwrap();
    writeln!(s, "    if (bid >= batch_size) return;").unwrap();
    writeln!(s, "    const unsigned int K = {k};").unwrap();
    writeln!(
        s,
        "    unsigned long long j_base = (unsigned long long)bid * num_variables * K;"
    )
    .unwrap();

    // Initialize
    writeln!(s, "    for (unsigned int i = 0; i < num_variables; i++) {{").unwrap();
    writeln!(s, "        unsigned long long off = j_base + i * K;").unwrap();
    writeln!(s, "        jets[off] = constants[i];").unwrap();
    for c in 1..k {
        writeln!(s, "        jets[off + {c}] = (F)0;").unwrap();
    }
    writeln!(s, "    }}").unwrap();

    // Set inputs
    writeln!(
        s,
        "    unsigned long long in_base = (unsigned long long)bid * num_inputs;"
    )
    .unwrap();
    writeln!(s, "    for (unsigned int i = 0; i < num_inputs; i++) {{").unwrap();
    writeln!(s, "        unsigned long long off = j_base + i * K;").unwrap();
    writeln!(s, "        jets[off] = primal_inputs[in_base + i];").unwrap();
    if k > 1 {
        writeln!(s, "        jets[off + 1] = direction_seeds[in_base + i];").unwrap();
    }
    writeln!(s, "    }}").unwrap();

    // Tape walk
    writeln!(
        s,
        "    for (unsigned int i = num_inputs; i < num_ops; i++) {{"
    )
    .unwrap();
    writeln!(s, "        unsigned int op = opcodes[i];").unwrap();
    writeln!(s, "        if (op == OP_CONST) continue;").unwrap();
    writeln!(s, "        unsigned int a_idx = arg0[i];").unwrap();
    writeln!(s, "        unsigned int b_idx = arg1[i];").unwrap();

    // Load a
    writeln!(s, "        JetK a;").unwrap();
    writeln!(
        s,
        "        unsigned long long a_off = j_base + (unsigned long long)a_idx * K;"
    )
    .unwrap();
    for c in 0..k {
        writeln!(s, "        a.v[{c}] = jets[a_off + {c}];").unwrap();
    }
    writeln!(s, "        JetK r;").unwrap();

    // Switch on opcode
    writeln!(s, "        switch (op) {{").unwrap();

    // Binary ops
    for (case, func) in &[
        (2, "jet_add"),
        (3, "jet_sub"),
        (4, "jet_mul"),
        (5, "jet_div"),
    ] {
        writeln!(s, "        case {case}: {{").unwrap();
        writeln!(
            s,
            "            JetK b; unsigned long long b_off = j_base + (unsigned long long)b_idx * K;"
        )
        .unwrap();
        for c in 0..k {
            writeln!(s, "            b.v[{c}] = jets[b_off + {c}];").unwrap();
        }
        writeln!(s, "            r = {func}(a, b); break;").unwrap();
        writeln!(s, "        }}").unwrap();
    }

    // REM: a % b = a - trunc(a[0]/b[0]) * b. Higher-order: r[k] = a[k] - q * b[k].
    writeln!(s, "        case 6: {{").unwrap();
    writeln!(
        s,
        "            JetK b; unsigned long long b_off = j_base + (unsigned long long)b_idx * K;"
    )
    .unwrap();
    for c in 0..k {
        writeln!(s, "            b.v[{c}] = jets[b_off + {c}];").unwrap();
    }
    // Zero-leading divisor → all-NaN jet (matches CPU Taylor::rem).
    writeln!(s, "            if (b.v[0] == (F)0) {{").unwrap();
    for i in 0..k {
        writeln!(s, "                r.v[{i}] = F(0.0/0.0);").unwrap();
    }
    writeln!(s, "                break;").unwrap();
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            F q = trunc(a.v[0] / b.v[0]);").unwrap();
    writeln!(s, "            r.v[0] = fmod(a.v[0], b.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "            r.v[{i}] = a.v[{i}] - q * b.v[{i}];").unwrap();
    }
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // POWF — a constant integer exponent (in i32 range) is routed to the powi
    // jet, finite for any base (a<0 is well-defined at integer powers),
    // mirroring CPU `taylor_powf`. A non-integer/variable exponent at a<=0 is NaN.
    writeln!(s, "        case 7: {{").unwrap();
    writeln!(
        s,
        "            JetK b; unsigned long long b_off = j_base + (unsigned long long)b_idx * K;"
    )
    .unwrap();
    for c in 0..k {
        writeln!(s, "            b.v[{c}] = jets[b_off + {c}];").unwrap();
    }
    // `< 2^31` (not `<= i32::MAX`): 2147483647 rounds up to 2^31 in f32, so the
    // `<=` form would admit 2^31 and overflow the `(int)` cast below.
    let mut b_int_cond = String::from("b.v[0] == round(b.v[0]) && fabs(b.v[0]) < F(2147483648.0)");
    for i in 1..k {
        b_int_cond.push_str(&format!(" && b.v[{i}] == F(0)"));
    }
    writeln!(s, "            bool b_is_int = {b_int_cond};").unwrap();
    writeln!(s, "            if (b_is_int) {{").unwrap();
    write_cuda_powi_jet(s, k, "b.v[0]", "(int)b.v[0]");
    // Negative base with a non-constant-integer exponent → all-NaN jet (matches
    // CPU `taylor_powf` and the WGSL twin). C `pow(-2,3) = -8` is finite, so
    // force NaN here — otherwise the primal would be finite beside NaN
    // derivative coefficients (a mixed jet), diverging from CPU.
    writeln!(s, "            }} else if (a.v[0] < F(0)) {{").unwrap();
    for i in 0..k {
        writeln!(s, "                r.v[{i}] = F(0.0/0.0);").unwrap();
    }
    writeln!(s, "            }} else if (a.v[0] <= F(0)) {{").unwrap();
    writeln!(s, "                F val = pow(a.v[0], b.v[0]);").unwrap();
    writeln!(
        s,
        "                F da = b.v[0] * pow(a.v[0], b.v[0] - F(1));"
    )
    .unwrap();
    // Reached only for a non-integer or variable exponent at a <= 0, where
    // a^b is genuinely NaN — the constant-integer case is routed to the powi
    // jet above. db = 0 matches CPU OpCode::Powf.
    writeln!(s, "                F db = F(0);").unwrap();
    writeln!(s, "                r.v[0] = val;").unwrap();
    if k > 1 {
        writeln!(s, "                r.v[1] = da * a.v[1] + db * b.v[1];").unwrap();
    }
    for i in 2..k {
        // CPU powf with negative base produces NaN through exp(b*ln(negative)).
        // Use explicit NaN to match CPU behavior regardless of whether val is finite.
        writeln!(s, "                r.v[{i}] = F(0.0/0.0);").unwrap();
    }
    writeln!(s, "            }} else {{").unwrap();
    writeln!(s, "                JetK lna = jet_ln(a);").unwrap();
    writeln!(s, "                JetK product = jet_mul(b, lna);").unwrap();
    writeln!(s, "                r = jet_exp(product);").unwrap();
    writeln!(s, "                r.v[0] = pow(a.v[0], b.v[0]);").unwrap();
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // ATAN2 — guard b==0 (division by zero in jet_div)
    writeln!(s, "        case 8: {{").unwrap();
    writeln!(
        s,
        "            JetK b; unsigned long long b_off = j_base + (unsigned long long)b_idx * K;"
    )
    .unwrap();
    for c in 0..k {
        writeln!(s, "            b.v[{c}] = jets[b_off + {c}];").unwrap();
    }
    writeln!(s, "            if (b.v[0] == F(0)) {{").unwrap();
    writeln!(s, "                r.v[0] = atan2(a.v[0], b.v[0]);").unwrap();
    writeln!(s, "                if (a.v[0] == F(0)) {{").unwrap();
    for i in 1..k {
        writeln!(s, "                    r.v[{i}] = F(0);").unwrap();
    }
    // b.v[0]==0, a.v[0]!=0: use atan2(a,b) = -atan(b/a) + pi/2*sign(a).
    // Compute full Taylor composition: ratio=b/a (ratio.v[0]=0), c=atan(ratio), negate.
    writeln!(s, "                }} else {{").unwrap();
    writeln!(s, "                    JetK ratio = jet_div(b, a);").unwrap();
    writeln!(s, "                    JetK c = jet_atan(ratio);").unwrap();
    for i in 1..k {
        writeln!(s, "                    r.v[{i}] = -c.v[{i}];").unwrap();
    }
    writeln!(s, "                }}").unwrap();
    writeln!(s, "            }} else {{").unwrap();
    writeln!(s, "                r = jet_atan(jet_div(a, b));").unwrap();
    writeln!(s, "                r.v[0] = atan2(a.v[0], b.v[0]);").unwrap();
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // HYPOT. Full jet-wide max-rescale to keep every coefficient
    // bounded even when `a.v[0] ~ 1e20`. Mirrors CPU
    // `taylor_ops::taylor_hypot` and the WGSL emitter above. Branch
    // order: Inf first, then NaN, then `h == 0` (with the one-level
    // shift-and-square unroll), then general — preserves IEEE
    // `hypot(Inf, NaN) = Inf` and `hypot(NaN, 0) = NaN`.
    //
    // Function-domain-boundary conventions (matching CPU):
    // - `hypot(Inf, finite)` — primal = Inf, higher-order = NaN
    //   (CPU's rescale path produces `Inf * 0 = NaN`).
    // - `hypot(0, 0)` with `a.v[1] != 0 || b.v[1] != 0` — one-level
    //   shift-and-square unroll: rescale+sqrt on the shifted jet,
    //   then shift result right by 1.
    // - `hypot(0, 0)` with all higher-order seeds also zero —
    //   0 primal + Inf higher (matches CPU's `taylor_sqrt` at a
    //   zero leading coefficient).
    writeln!(s, "        case 9: {{").unwrap();
    writeln!(
        s,
        "            JetK b; unsigned long long b_off = j_base + (unsigned long long)b_idx * K;"
    )
    .unwrap();
    for c in 0..k {
        writeln!(s, "            b.v[{c}] = jets[b_off + {c}];").unwrap();
    }
    writeln!(s, "            F aa = fabs(a.v[0]);").unwrap();
    writeln!(s, "            F bb = fabs(b.v[0]);").unwrap();
    // CUDA: `isinf` is the correct way to detect ±Inf for both f32 and f64.
    writeln!(s, "            if (isinf(aa) || isinf(bb)) {{").unwrap();
    // Inf branch: NVRTC's default program doesn't expose the
    // `INFINITY` macro, so synthesise +Inf via `fmax` (at least one
    // of {aa, bb} is +Inf and both are non-negative). Higher-order
    // emits NaN via `inf - inf` (runtime) to match CPU's
    // `Inf * 0 = NaN` rescale-path output.
    writeln!(s, "                F inf_sentinel = fmax(aa, bb);").unwrap();
    writeln!(
        s,
        "                F nan_sentinel = inf_sentinel - inf_sentinel;"
    )
    .unwrap();
    writeln!(s, "                r.v[0] = inf_sentinel;").unwrap();
    for i in 1..k {
        writeln!(s, "                r.v[{i}] = nan_sentinel;").unwrap();
    }
    writeln!(s, "            }} else if (isnan(aa) || isnan(bb)) {{").unwrap();
    // NaN branch: without this, `fmax(NaN, 0) = 0` (IEEE 754-2008
    // maxNum returns the non-NaN argument), causing `hypot(NaN, 0)`
    // to enter the `h == 0` branch and silently return the zero jet
    // — swallowing the NaN that IEEE `hypot` requires propagating.
    // `a.v[0] + b.v[0]` is the simplest NaN-propagator (NaN + anything
    // = NaN; both NaN: still NaN).
    writeln!(s, "                r.v[0] = a.v[0] + b.v[0];").unwrap();
    for i in 1..k {
        writeln!(s, "                r.v[{i}] = (F)0;").unwrap();
    }
    writeln!(s, "            }} else {{").unwrap();
    writeln!(s, "                F h = fmax(aa, bb);").unwrap();
    // Zero branch (shift-and-square unroll).
    writeln!(s, "                if (h == (F)0) {{").unwrap();
    if k >= 2 {
        writeln!(s, "                    JetK a_shifted;").unwrap();
        writeln!(s, "                    JetK b_shifted;").unwrap();
        for i in 0..(k - 1) {
            writeln!(
                s,
                "                    a_shifted.v[{i}] = a.v[{ip1}];",
                ip1 = i + 1
            )
            .unwrap();
            writeln!(
                s,
                "                    b_shifted.v[{i}] = b.v[{ip1}];",
                ip1 = i + 1
            )
            .unwrap();
        }
        writeln!(
            s,
            "                    a_shifted.v[{last}] = (F)0;",
            last = k - 1
        )
        .unwrap();
        writeln!(
            s,
            "                    b_shifted.v[{last}] = (F)0;",
            last = k - 1
        )
        .unwrap();
        writeln!(
            s,
            "                    F h_inner = fmax(fabs(a_shifted.v[0]), fabs(b_shifted.v[0]));"
        )
        .unwrap();
        writeln!(s, "                    if (h_inner == (F)0) {{").unwrap();
        // Deeper-order-zero: 0 primal + Inf higher. +Inf via (F)1/(F)0
        // — NVRTC default rounding produces +Inf, not NaN. Using a
        // runtime divide rather than a bitcast keeps us portable
        // across NVRTC flag configurations.
        writeln!(s, "                        F pos_inf = (F)1 / (F)0;").unwrap();
        writeln!(s, "                        r.v[0] = (F)0;").unwrap();
        for i in 1..k {
            writeln!(s, "                        r.v[{i}] = pos_inf;").unwrap();
        }
        writeln!(s, "                    }} else {{").unwrap();
        writeln!(s, "                        F inv_h_inner = (F)1 / h_inner;").unwrap();
        writeln!(
            s,
            "                        JetK a_ss = jet_scale(a_shifted, inv_h_inner);"
        )
        .unwrap();
        writeln!(
            s,
            "                        JetK b_ss = jet_scale(b_shifted, inv_h_inner);"
        )
        .unwrap();
        writeln!(
            s,
            "                        JetK sum_sq_s = jet_add(jet_mul(a_ss, a_ss), jet_mul(b_ss, b_ss));"
        )
        .unwrap();
        writeln!(
            s,
            "                        JetK r_s_inner = jet_sqrt(sum_sq_s);"
        )
        .unwrap();
        writeln!(
            s,
            "                        JetK inner = jet_scale(r_s_inner, h_inner);"
        )
        .unwrap();
        writeln!(s, "                        r.v[0] = (F)0;").unwrap();
        for i in 1..k {
            writeln!(
                s,
                "                        r.v[{i}] = inner.v[{im1}];",
                im1 = i - 1
            )
            .unwrap();
        }
        writeln!(s, "                    }}").unwrap();
    } else {
        // K == 1: no higher-order slots. Primal-only zero result.
        writeln!(s, "                    r.v[0] = (F)0;").unwrap();
    }
    writeln!(s, "                }} else {{").unwrap();
    writeln!(s, "                    F inv_h = (F)1 / h;").unwrap();
    writeln!(s, "                    JetK a_s = jet_scale(a, inv_h);").unwrap();
    writeln!(s, "                    JetK b_s = jet_scale(b, inv_h);").unwrap();
    writeln!(
        s,
        "                    JetK sum_sq = jet_add(jet_mul(a_s, a_s), jet_mul(b_s, b_s));"
    )
    .unwrap();
    writeln!(s, "                    JetK r_s = jet_sqrt(sum_sq);").unwrap();
    writeln!(s, "                    r = jet_scale(r_s, h);").unwrap();
    // CUDA: `hypot()` is a load-bearing primal override — NVRTC's
    // intrinsic rounds per-device and may differ by a ULP from the
    // explicit rescale formula even at moderate magnitudes.
    writeln!(s, "                    r.v[0] = hypot(a.v[0], b.v[0]);").unwrap();
    writeln!(s, "                }}").unwrap();
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // MAX, MIN — `|| isnan(b.v[0])` mirrors CPU Taylor::max/min's `|| other.is_nan()`
    // tie-break, so max(x, NaN) returns the finite operand instead of the NaN.
    for (case, cmp) in &[(10, ">="), (11, "<=")] {
        writeln!(s, "        case {case}: {{").unwrap();
        writeln!(
            s,
            "            JetK b; unsigned long long b_off = j_base + (unsigned long long)b_idx * K;"
        )
        .unwrap();
        for c in 0..k {
            writeln!(s, "            b.v[{c}] = jets[b_off + {c}];").unwrap();
        }
        writeln!(
            s,
            "            r = (a.v[0] {cmp} b.v[0] || isnan(b.v[0])) ? a : b; break;"
        )
        .unwrap();
        writeln!(s, "        }}").unwrap();
    }

    // Unary
    writeln!(s, "        case 12: r = jet_neg(a); break;").unwrap();
    writeln!(s, "        case 13: r = jet_recip(a); break;").unwrap();
    writeln!(s, "        case 14: r = jet_sqrt(a); break;").unwrap();

    // CBRT — guard a.v[0]==0 (vertical tangent singularity)
    writeln!(s, "        case 15: {{").unwrap();
    writeln!(s, "            if (a.v[0] == F(0)) {{").unwrap();
    writeln!(s, "                r.v[0] = F(0);").unwrap();
    for i in 1..k {
        writeln!(s, "                r.v[{i}] = F(1.0/0.0);").unwrap();
    }
    writeln!(s, "            }} else {{").unwrap();
    writeln!(s, "                F sg = _sign(a.v[0]);").unwrap();
    writeln!(s, "                JetK abs_a; abs_a.v[0] = fabs(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "                abs_a.v[{i}] = sg * a.v[{i}];").unwrap();
    }
    writeln!(
        s,
        "                JetK e = jet_exp(jet_scale(jet_ln(abs_a), (F)(1.0/3.0)));"
    )
    .unwrap();
    for i in 0..k {
        writeln!(s, "                r.v[{i}] = sg * e.v[{i}];").unwrap();
    }
    writeln!(s, "            }}").unwrap();
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // POWI — handle negative bases via sign(a)^n * exp(n * ln(|a|))
    writeln!(s, "        case 16: {{").unwrap();
    write_cuda_powi_jet(s, k, "(F)ni", "(int)b_idx");
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // Transcendental
    writeln!(s, "        case 17: r = jet_exp(a); break;").unwrap();
    writeln!(s, "        case 18: {{ r = jet_exp(jet_scale(a, log((F)2))); r.v[0] = exp2(a.v[0]); break; }}").unwrap();
    writeln!(
        s,
        "        case 19: {{ r = jet_exp(a); r.v[0] = expm1(a.v[0]); break; }}"
    )
    .unwrap();
    writeln!(s, "        case 20: r = jet_ln(a); break;").unwrap();

    // LOG2
    writeln!(s, "        case 21: {{").unwrap();
    writeln!(s, "            r = jet_ln(a);").unwrap();
    writeln!(s, "            F inv_ln2 = (F)1 / log((F)2);").unwrap();
    writeln!(s, "            r.v[0] = log2(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "            r.v[{i}] *= inv_ln2;").unwrap();
    }
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // LOG10
    writeln!(s, "        case 22: {{").unwrap();
    writeln!(s, "            r = jet_ln(a);").unwrap();
    writeln!(s, "            F inv_ln10 = (F)1 / log((F)10);").unwrap();
    writeln!(s, "            r.v[0] = log(a.v[0]) * inv_ln10;").unwrap();
    for i in 1..k {
        writeln!(s, "            r.v[{i}] *= inv_ln10;").unwrap();
    }
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // LN1P
    writeln!(s, "        case 23: {{").unwrap();
    writeln!(s, "            JetK opa; opa.v[0] = (F)1 + a.v[0];").unwrap();
    for i in 1..k {
        writeln!(s, "            opa.v[{i}] = a.v[{i}];").unwrap();
    }
    writeln!(
        s,
        "            r = jet_ln(opa); r.v[0] = log1p(a.v[0]); break;"
    )
    .unwrap();
    writeln!(s, "        }}").unwrap();

    // Sin, Cos, Tan, Asin, Acos, Atan, Sinh, Cosh, Tanh, Asinh, Acosh, Atanh
    writeln!(
        s,
        "        case 24: {{ JetPair sc = jet_sin_cos(a); r = sc.a; break; }}"
    )
    .unwrap();
    writeln!(
        s,
        "        case 25: {{ JetPair sc = jet_sin_cos(a); r = sc.b; break; }}"
    )
    .unwrap();
    writeln!(s, "        case 26: r = jet_tan(a); break;").unwrap();
    writeln!(s, "        case 27: r = jet_asin(a); break;").unwrap();
    writeln!(s, "        case 28: r = jet_acos(a); break;").unwrap();
    writeln!(s, "        case 29: r = jet_atan(a); break;").unwrap();
    writeln!(
        s,
        "        case 30: {{ JetPair sc = jet_sinh_cosh(a); r = sc.a; break; }}"
    )
    .unwrap();
    writeln!(
        s,
        "        case 31: {{ JetPair sc = jet_sinh_cosh(a); r = sc.b; break; }}"
    )
    .unwrap();
    writeln!(s, "        case 32: r = jet_tanh(a); break;").unwrap();
    writeln!(s, "        case 33: r = jet_asinh(a); break;").unwrap();
    writeln!(s, "        case 34: r = jet_acosh(a); break;").unwrap();
    writeln!(s, "        case 35: r = jet_atanh(a); break;").unwrap();

    // ABS
    // Nonsmooth — ABS: scan for the first non-zero Taylor coefficient to
    // decide the sign. Matches CPU `Taylor::abs`: at a = [0, -v, …] the
    // sign is -1 (driven by the first-non-zero coefficient), not +1.
    writeln!(s, "        case 36: {{").unwrap();
    writeln!(s, "            F sg = F(1);").unwrap();
    for i in 0..k {
        writeln!(
            s,
            "            if (a.v[{i}] != F(0)) {{ sg = (a.v[{i}] < F(0)) ? F(-1) : F(1); }} else",
        )
        .unwrap();
    }
    writeln!(
        s,
        "            {{ /* identically zero series → sg = 1 */ }}"
    )
    .unwrap();
    writeln!(s, "            r.v[0] = fabs(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "            r.v[{i}] = sg * a.v[{i}];").unwrap();
    }
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // SIGNUM..TRUNC
    writeln!(s, "        case 37: r = jet_const(_sign(a.v[0])); break;").unwrap();
    writeln!(s, "        case 38: r = jet_const(floor(a.v[0])); break;").unwrap();
    writeln!(s, "        case 39: r = jet_const(ceil(a.v[0])); break;").unwrap();
    writeln!(s, "        case 40: r = jet_const(round(a.v[0])); break;").unwrap();
    writeln!(s, "        case 41: r = jet_const(trunc(a.v[0])); break;").unwrap();

    // FRACT
    writeln!(s, "        case 42: {{").unwrap();
    writeln!(s, "            r.v[0] = _fract(a.v[0]);").unwrap();
    for i in 1..k {
        writeln!(s, "            r.v[{i}] = a.v[{i}];").unwrap();
    }
    writeln!(s, "            break;").unwrap();
    writeln!(s, "        }}").unwrap();

    writeln!(s, "        default: break;").unwrap();
    writeln!(s, "        }}").unwrap();

    // Store result
    writeln!(s, "        unsigned long long r_off = j_base + i * K;").unwrap();
    for c in 0..k {
        writeln!(s, "        jets[r_off + {c}] = r.v[{c}];").unwrap();
    }
    writeln!(s, "    }}").unwrap();

    // Write outputs
    writeln!(
        s,
        "    unsigned long long out_base = (unsigned long long)bid * num_outputs * K;"
    )
    .unwrap();
    writeln!(s, "    for (unsigned int j = 0; j < num_outputs; j++) {{").unwrap();
    writeln!(s, "        unsigned int oi = output_indices[j];").unwrap();
    writeln!(
        s,
        "        unsigned long long src = j_base + (unsigned long long)oi * K;"
    )
    .unwrap();
    writeln!(s, "        unsigned long long dst = out_base + j * K;").unwrap();
    for c in 0..k {
        writeln!(s, "        jet_outputs[dst + {c}] = jets[src + {c}];").unwrap();
    }
    writeln!(s, "    }}").unwrap();
    writeln!(s, "}}").unwrap();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn wgsl_k3_compiles() {
        let shader = generate_taylor_wgsl(3);
        assert!(shader.contains("struct JetK { v: array<f32, 3>, }"));
        assert!(shader.contains("fn jet_mul"));
        assert!(shader.contains("fn jet_exp"));
        assert!(shader.contains("fn main"));
    }

    #[test]
    fn cuda_k3_compiles() {
        let kernel = generate_taylor_cuda(3);
        assert!(kernel.contains("F v[3]"));
        assert!(kernel.contains("jet_mul"));
        assert!(kernel.contains("jet_exp"));
        assert!(kernel.contains("taylor_forward_kth"));
    }

    #[test]
    fn all_k_values_generate() {
        for k in 1..=5 {
            let wgsl = generate_taylor_wgsl(k);
            assert!(wgsl.contains(&format!("array<f32, {k}>")));
            let cuda = generate_taylor_cuda(k);
            assert!(cuda.contains(&format!("F v[{k}]")));
        }
    }

    // asin/acos/atanh jet denominators use the cancellation-safe factored
    // `(1 - a0)(1 + a0)` (matching CPU `taylor_asin`/`taylor_atanh` and the
    // already-factored `jet_acosh`) rather than the naive `1 - a0^2`, which
    // loses the `eps^2` term near `|a0| = 1`. The higher coefficients stay
    // `-asq.v[i]` (bit-identical to the factored expansion). The f32 GPU path
    // cannot demonstrate the precision gain within the parity tolerance floor
    // (see gpu_cpu_parity.rs acosh near-1 note), so this codegen assertion is
    // the load-bearing regression pin.
    #[test]
    fn m7_inverse_trig_denominators_factored() {
        let wgsl = generate_taylor_wgsl(3);
        // asin/acos/atanh each emit the factored leading denominator.
        assert_eq!(
            wgsl.matches("d.v[0] = (1.0 - a.v[0]) * (1.0 + a.v[0]);")
                .count(),
            3,
            "asin/acos/atanh must all emit the factored `(1-a0)(1+a0)` denominator"
        );
        // The naive unfactored form must be gone (atan keeps `1.0 + asq.v[0]`).
        assert!(
            !wgsl.contains("d.v[0] = 1.0 - asq.v[0];"),
            "no jet denominator may keep the naive `1 - a0^2` form"
        );

        let cuda = generate_taylor_cuda(3);
        assert!(
            cuda.contains("d.v[0] = ((F)1 - a.v[0]) * ((F)1 + a.v[0]);"),
            "CUDA `1.0 -` table branch must emit the factored denominator"
        );
        assert!(
            !cuda.contains("d.v[0] = (F)1 - asq.v[0];"),
            "CUDA must not keep the naive `1 - a0^2` denominator"
        );
    }

    // EXPM1/LN1P patch only the primal `v[0]`; the naive `exp(x)-1` / `log(1+x)`
    // lose accuracy for small `x`. WGSL uses guard-free polyfills: a Taylor/Mercator
    // series for `|x| < 0.125` (no cancellation) and the naive `exp(x)-1.0` /
    // `log(1.0+x)` above the threshold, which give the correct limits directly
    // (`+Inf` for large x since `+Inf - 1 = +Inf`; `-Inf` at `x = -1` since
    // `log(0) = -Inf`). CUDA uses the native `expm1`/`log1p` device intrinsics.
    #[test]
    fn m9_expm1_ln1p_accurate_primal() {
        let wgsl = generate_taylor_wgsl(3);
        assert!(wgsl.contains("fn expm1_f(x: f32) -> f32"));
        assert!(wgsl.contains("fn log1p_f(x: f32) -> f32"));
        assert!(wgsl.contains("r.v[0] = expm1_f(a.v[0]);"));
        assert!(wgsl.contains("r.v[0] = log1p_f(a.v[0]);"));
        // Both polyfills use the series-below-0.125 / naive-above split (the Kahan
        // trick is defeated by Metal fast-math: it reassociates expm1's grouping to
        // overflow, and folds log1p's `(1+x)-1` back to `x`). expm1_f falls back to
        // `exp(x) - 1.0` above the threshold (right limits: +Inf / -1 without guards).
        let expm1_body = wgsl[wgsl.find("fn expm1_f").unwrap()..]
            .split("fn log1p_f")
            .next()
            .unwrap();
        assert!(
            expm1_body.contains("abs(x) < 0.125") && expm1_body.contains("exp(x) - 1.0"),
            "expm1_f must use the series-below-0.125 / naive-above split"
        );
        assert!(
            wgsl[wgsl.find("fn log1p_f").unwrap()..].contains("abs(x) < 0.125"),
            "log1p_f must use the series-below-0.125 split"
        );
        // The naive primal forms must be gone.
        assert!(!wgsl.contains("r.v[0] = exp(a.v[0]) - 1.0;"));
        assert!(!wgsl.contains("r.v[0] = log(1.0 + a.v[0]);"));

        let cuda = generate_taylor_cuda(3);
        assert!(cuda.contains("r.v[0] = expm1(a.v[0]);"));
        assert!(cuda.contains("r.v[0] = log1p(a.v[0]);"));
        assert!(!cuda.contains("r.v[0] = exp(a.v[0]) - (F)1;"));
        assert!(!cuda.contains("r.v[0] = log((F)1 + a.v[0]);"));
    }

    // Series domain-NaN guards: jet_ln (covers ln/log2/log10/ln1p by delegation),
    // jet_acosh, and jet_atanh emit an all-NaN jet strictly outside the real domain,
    // mirroring the CPU taylor_ln/taylor_acosh/taylor_atanh guards. Strict `<`
    // comparisons keep the branch-point boundary (ln@0, acosh@1, atanh@±1) singular.
    #[test]
    fn m_series_domain_nan_guards_present() {
        let wgsl = generate_taylor_wgsl(3);
        assert!(wgsl.contains("if (a.v[0] < 0.0) {"), "jet_ln domain guard");
        assert!(
            wgsl.contains("if (a.v[0] < 1.0) {"),
            "jet_acosh domain guard"
        );
        assert!(
            wgsl.contains("if (a.v[0] < -1.0 || a.v[0] > 1.0) {"),
            "jet_atanh domain guard"
        );

        let cuda = generate_taylor_cuda(3);
        assert!(
            cuda.contains("if (a.v[0] < (F)0) {"),
            "CUDA jet_ln domain guard"
        );
        assert!(
            cuda.contains("if (a.v[0] < (F)1) {"),
            "CUDA jet_acosh domain guard"
        );
        assert!(
            cuda.contains("if (a.v[0] < (F)-1 || a.v[0] > (F)1) {"),
            "CUDA jet_atanh domain guard"
        );
    }

    #[test]
    fn m_powf_neg_base_all_nan_in_both_backends() {
        // taylor_powf convention (Option A): a negative base with a
        // non-constant-integer exponent gives an all-NaN jet on both backends.
        // CUDA's C `pow(-2,3) = -8` is forced to NaN to match CPU + WGSL,
        // avoiding a mixed (finite primal / NaN derivative) jet.
        let wgsl = generate_taylor_wgsl(3);
        assert!(
            wgsl.contains("} else if a.v[0] < 0.0 {"),
            "WGSL powf neg-base all-NaN branch"
        );
        let cuda = generate_taylor_cuda(3);
        assert!(
            cuda.contains("} else if (a.v[0] < F(0)) {"),
            "CUDA powf neg-base all-NaN branch"
        );
    }

    #[test]
    fn jet_ln_omits_inv_a0_at_k1() {
        // At K=1 the ln recurrence loop is empty, so `inv_a0` must not be
        // emitted — an unused binding is a shader-compile warning.
        assert!(
            !generate_taylor_wgsl(1).contains("inv_a0"),
            "K=1 WGSL jet_ln must not emit inv_a0"
        );
        assert!(
            generate_taylor_wgsl(3).contains("inv_a0"),
            "K=3 WGSL jet_ln must emit inv_a0"
        );
        assert!(
            !generate_taylor_cuda(1).contains("inv_a0"),
            "K=1 CUDA jet_ln must not emit inv_a0"
        );
        assert!(
            generate_taylor_cuda(3).contains("inv_a0"),
            "K=3 CUDA jet_ln must emit inv_a0"
        );
    }

    // GPU-low parity: MAX/MIN keep CPU's NaN tie-break (`|| other.is_nan()`), jet_sqrt
    // gives a uniform +Inf higher-coeff jet at a0==0 (not sign-dependent Inf/NaN), and
    // REM by a zero-leading divisor returns an all-NaN jet (matching CPU Taylor::rem).
    #[test]
    fn m_gpu_low_parity_max_min_sqrt_rem() {
        let wgsl = generate_taylor_wgsl(3);
        // MAX/MIN: the NaN-of-b tie-break makes max(x, NaN) return the finite operand.
        assert_eq!(
            wgsl.matches("|| is_nan_f32(b.v[0])").count(),
            2,
            "WGSL MAX and MIN must both carry the NaN tie-break"
        );
        // jet_sqrt: uniform +Inf branch at a0==0.
        assert!(
            wgsl.contains("fn jet_sqrt") && wgsl.contains("if (a.v[0] == 0.0) {"),
            "WGSL jet_sqrt must guard a0==0"
        );
        // REM: all-NaN when the divisor's leading coeff is zero.
        assert!(
            wgsl.contains("if (b.v[0] == 0.0) {"),
            "WGSL REM must guard a zero-leading divisor"
        );

        let cuda = generate_taylor_cuda(3);
        assert!(
            cuda.contains("|| isnan(b.v[0])"),
            "CUDA MAX/MIN must carry the NaN tie-break"
        );
        assert!(
            cuda.contains("if (a.v[0] == (F)0) {"),
            "CUDA jet_sqrt must guard a0==0"
        );
        assert!(
            cuda.contains("if (b.v[0] == (F)0) {"),
            "CUDA REM must guard a zero-leading divisor"
        );
    }

    // GPU-low (WGSL primal conventions): SIGNUM uses the sign-bit helper (so -0.0
    // gives -1, not +1), and ROUND is ties-away-from-zero (WGSL `round()` is
    // ties-to-even). CUDA's `round`/`copysign` are already ties-away / sign-bit.
    #[test]
    fn m_wgsl_signum_round_conventions() {
        let wgsl = generate_taylor_wgsl(3);
        assert!(wgsl.contains("fn signum_f32"));
        assert!(wgsl.contains("val = signum_f32(a.v[0]);"));
        assert!(!wgsl.contains("val = select(sign(a.v[0]), 1.0"));
        // Ties-away via trunc-then-fractional-compare (`a - trunc(a)` is exact, so
        // no double-rounding — the additive `trunc(a+0.5)` form regresses large
        // odd integers and `nextbelow(0.5)`).
        assert!(wgsl.contains(
            "let t = trunc(a.v[0]); val = select(t, t + select(-1.0, 1.0, a.v[0] >= 0.0), abs(a.v[0] - t) >= 0.5);"
        ));
        assert!(!wgsl.contains("val = round(a.v[0]);"));
        assert!(!wgsl.contains("trunc(a.v[0] + select(-0.5, 0.5"));
    }
}
