pub const ADAM: &str = r#"
#include <metal_stdlib>
using namespace metal;
kernel void optirs_adam(
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];
if (wd > 0.0f) { g = g + wd * p; }
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;
x[idx] = p - lr * m_hat / (sqrt(v_hat) + eps);
}
"#;Expand description
Adam with coupled L2 weight decay. See super::wgsl::ADAM for bindings.