Example

mlp_forward_static.kina

Source

let mlp = fn(x: tensor<F32, [128, 256], @host>, w1: tensor<F32, [256, 512], @host>, w2: tensor<F32, [512, 10], @host>) -> tensor<F32, [128, 10], @host> {
    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], @host> { 1.0 } in
let w1 = tensor<F32, [256, 512], @host> { 1.0 } in
let w2 = tensor<F32, [512, 10], @host> { 1.0 } in
mlp(x, w1, w2)