Example

ad_mlp.kina

Source

let mlp_forward = fn(x: tensor<F32, [2, 2], @host>, w1: tensor<F32, [2, 2], @host>, w2: tensor<F32, [2, 2], @host>) !{Diff} -> tensor<F32, [2, 2], @host> {
  let active_w1 = perform Diff.active(w1) in
  let active_w2 = perform Diff.active(w2) in
  let hidden = x @ active_w1 in
  let hidden_sq = hidden * hidden in
  let out = hidden_sq @ active_w2 in
  out
} in

let run_mlp = fn(x: tensor<F32, [2, 2], @host>, w1: tensor<F32, [2, 2], @host>, w2: tensor<F32, [2, 2], @host>) -> tensor<F32, [2, 2], @host> {
  handle mlp_forward(x, w1, w2) with {
    return(res) => res,
    Diff.active(v), k => k(v)
  }
} in
let x = tensor<F32, [2, 2], @host> { 1.0, 1.0, 1.0, 1.0 } in
let w1 = tensor<F32, [2, 2], @host> { 1.0, 1.0, 1.0, 1.0 } in
let w2 = tensor<F32, [2, 2], @host> { 1.0, 1.0, 1.0, 1.0 } in
run_mlp(x, w1, w2)