Constant ADAMW 

pub const ADAMW: &str = r#"
#include <metal_stdlib>
using namespace metal;

kernel void optirs_adamw(
    device float* x [[buffer(0)]],
    device const float* y [[buffer(1)]],
    device float* a [[buffer(2)]],
    device float* b [[buffer(3)]],
    device const float* result [[buffer(4)]],
    uint idx [[thread_position_in_grid]])
{
    uint n = as_type<uint>(result[7]);
    if (idx >= n) { return; }

    float lr = result[0];
    float beta1 = result[1];
    float beta2 = result[2];
    float eps = result[3];
    float wd = result[4];
    float bc1 = result[5];
    float bc2 = result[6];

    float p = x[idx];
    float g = y[idx];

    float mi = beta1 * a[idx] + (1.0f - beta1) * g;
    float vi = beta2 * b[idx] + (1.0f - beta2) * g * g;
    a[idx] = mi;
    b[idx] = vi;

    float m_hat = mi / bc1;
    float v_hat = vi / bc2;

    float decayed = p - lr * wd * p;
    x[idx] = decayed - lr * m_hat / (sqrt(v_hat) + eps);
}
"#;

AdamW with decoupled weight decay. See super::wgsl::ADAMW.