Autograd in Golang

autograd.go

package main

import "fmt"

type Value struct {
  data, grad float64
  backward   func()
  children   []*Value
}

func newValue(data float64, children []*Value) *Value {
  return &Value{data: data, children: children, backward: func() {}}
}

func (v *Value) String() string {
  return fmt.Sprintf("data: %f, grad: %f, children: %v", v.data, v.grad, v.children)
}

func (v *Value) add(v1 *Value) *Value {
  nv := newValue(v.data+v1.data, []*Value{v, v1})
  nv.backward = func() {
    v.grad += nv.grad
    v1.grad += nv.grad
  }
  return nv
}

func (v *Value) mul(v1 *Value) *Value {
  nv := newValue(v.data*v1.data, []*Value{v, v1})
  nv.backward = func() {
    v.grad += v1.data * nv.grad
    v1.grad += v.data * nv.grad
  }
  return nv
}

func backprop(v *Value) {
  topo, visited := []*Value{}, map[*Value]bool{}
  var buildTopo func(v1 *Value)
  buildTopo = func(v1 *Value) {
    if !visited[v1] {
      visited[v1] = true
      for _, child := range v1.children {
        buildTopo(child)
      }
      topo = append(topo, v1)
    }
  }
  buildTopo(v)
  v.grad = 1
  for i := len(topo) - 1; i >= 0; i-- {
    topo[i].backward()
  }
}

func main() {
  a, b, c := newValue(1, nil), newValue(2, nil), newValue(3, nil)
  o := a.mul(b).add(c).add(c)
  backprop(o)
  fmt.Printf("o.data=%f, a.grad=%f, b.grad=%f, c.grad=%f", o.data, a.grad, b.grad, c.grad)
}