Example

ad_stdlib_test.kina

Source

let x = tensor<F32, [1, 32, 64], @host> { 1.0 } in
let gamma = tensor<F32, [1, 1, 64], @host> { 1.0 } in

// Test RMSNorm Gradients
let out1 = handle rmsnorm(x, gamma) with reverseMode in

// Test SwiGLU Gradients
let w1 = tensor<F32, [1, 64, 128], @host> { 1.0 } in
let w2 = tensor<F32, [1, 128, 64], @host> { 1.0 } in
let out2 = handle swiglu(x, w1, w2) with reverseMode in

// Test Attention Gradients
let wq = tensor<F32, [1, 64, 64], @host> { 1.0 } in
let wk = tensor<F32, [1, 64, 64], @host> { 1.0 } in
let wv = tensor<F32, [1, 64, 64], @host> { 1.0 } in
let wo = tensor<F32, [1, 64, 64], @host> { 1.0 } in

handle attention(x, wq, wk, wv, wo) with reverseMode