Example

mlp_forward_gpu.kina

Source

let mlp_gpu = fn(x: tensor<F32, [128, 256], @device(0)>, w1: tensor<F32, [256, 512], @device(0)>, w2: tensor<F32, [512, 10], @device(0)>) -> tensor<F32, [128, 10], @device(0)> {
    let h1 = x @ w1 in
    let h1_relu = relu(h1) in
    let out = h1_relu @ w2 in
    out
} in
let x = tensor<F32, [128, 256], @device(0)> { 1.0 } in
let w1 = tensor<F32, [256, 512], @device(0)> { 1.0 } in
let w2 = tensor<F32, [512, 10], @device(0)> { 1.0 } in
mlp_gpu(x, w1, w2)