Micrograd

Mix.install([
  {:axon, "~> 0.6.0"},
  {:dg, "~> 0.4.0"},
  {:scholar, "~> 0.2.1"},
  {:exla, "~> 0.6.1"},
  {:kino_vega_lite, "~> 0.1.10"},
  {:kino, "~> 0.11.0"}
])

Nx.global_default_backend(EXLA.Backend)
alias VegaLite, as: Vl

Node of gradient calculation (aka Value in micrograd)

defmodule MicrogradNode do
  @moduledoc "stores a single scalar value and its gradient"
  defstruct data: nil,
            grad: 0,
            label: "",
            _backward: nil,
            _op: "",
            _children: []

  @type t :: %MicrogradNode{
          data: Integer.t() | nil,
          grad: Integer.t(),
          label: String.t(),
          _backward: fun(any),
          _op: String.t(),
          _children: list(%MicrogradNode{})
        }

  def add(%MicrogradNode{} = left, %MicrogradNode{} = right, label \\ "add", grad \\ 0) do
    out = %MicrogradNode{
      data: left.data + right.data,
      grad: grad,
      _children: previous(left, right),
      _op: "+",
      label: label
    }

    backward = fn node ->
      left = %MicrogradNode{left | grad: left.grad + node.grad}
      right = %MicrogradNode{right | grad: right.grad + node.grad}

      %MicrogradNode{node | _children: previous(left, right)}
    end

    %MicrogradNode{out | _backward: backward}
  end

  def mult(%MicrogradNode{} = left, %MicrogradNode{} = right, label \\ "mult", grad \\ 0) do
    out = %MicrogradNode{
      data: left.data * right.data,
      grad: grad,
      _children: previous(left, right),
      _op: "*",
      label: label
    }

    backward = fn node ->
      left = %MicrogradNode{left | grad: right.data * node.grad}
      right = %MicrogradNode{right | grad: left.data * node.grad}

      %MicrogradNode{node | _children: previous(left, right)}
    end

    %MicrogradNode{out | _backward: backward}
  end

  def tanh(%MicrogradNode{data: data} = children, label \\ "tanh", grad \\ 1) do
    t = (:math.exp(2 * data) - 1) / (:math.exp(2 * data) + 1)

    out = %MicrogradNode{
      data: t,
      grad: grad,
      label: label,
      _children: [children],
      _op: "tanh"
    }

    backward = fn node ->
      children = %MicrogradNode{children | grad: 1 - t ** 2}
      %MicrogradNode{node | _children: [children]}
    end

    %MicrogradNode{out | _backward: backward}
  end

  def previous(left, right) do
    [left, right]
  end

  def backward(%MicrogradNode{_children: []} = root) do
    root
  end

  def backward(
        %MicrogradNode{
          data: data,
          grad: grad,
          label: label,
          _op: op,
          _backward: backward
        } = node
      ) do
    node =
      if node._backward do
        node._backward.(node)
      else
        node
      end

    updated_children =
      Enum.map(node._children, fn child ->
        backward(child)
      end)

    %MicrogradNode{
      data: data,
      grad: grad,
      label: label,
      _backward: backward,
      _op: op,
      _children: updated_children
    }
  end
end

Test micrograd node API

a = %MicrogradNode{data: 2, label: "a"}
b = %MicrogradNode{data: -3, label: "b"}
MicrogradNode.add(a, b)

a = %MicrogradNode{data: 2, label: "a"}
b = %MicrogradNode{data: -3, label: "b"}
c = %MicrogradNode{data: 10, label: "c"}
d = MicrogradNode.mult(a, b, "e") |> MicrogradNode.add(c, "d", 1)

Visualization

defmodule Graph do
  def draw_dot(root) do
    dot = DG.new()
    build_dot(root, dot, %{count: 0, ops: [], rids: []})

    dot
  end

  def build_dot(node, dot, visited) do
    label = "#{node.label} -> data #{node.data} -> grad #{node.grad}"
    rid = to_string(:rand.uniform(1000))

    DG.add_vertex(dot, node.label, label)

    count = Map.get(visited, :count)
    ops = Map.get(visited, :ops)
    rids = Map.get(visited, :rids)
    val = [node._op]
    visited = %{count: count + 1, ops: ops ++ val, rids: rids ++ [rid]}

    if node._op != "" do
      # create left _op vertex and connect to right edge
      DG.add_vertex(dot, "#{node._op}" <> rid, node._op)
      DG.add_edge(dot, "#{node._op}" <> rid, node.label)
    end

    if count != 0 do
      # if not root node, create edge between vertex and _op
      ops = Map.get(visited, :ops)
      op = Enum.at(ops, count - 1)
      rid = Enum.at(rids, count - 1)
      DG.add_edge(dot, String.trim(node.label), String.trim(op) <> rid)
    end

    Enum.map(node._children, fn child ->
      build_dot(child, dot, visited)
    end)
  end
end

d
|> MicrogradNode.backward()
|> Graph.draw_dot()
|> inspect()
|> Kino.Mermaid.new()

Nx

x1 = Nx.tensor([2.0])
x2 = Nx.tensor([0.0])

w1 = Nx.tensor([-3.0])
w2 = Nx.tensor([1.0])

b = Nx.tensor([6.8813735870195432])

x1w1 = Nx.multiply(x1, w1)
x2w2 = Nx.multiply(x2, w2)

output = x1w1 |> Nx.add(x2w2) |> Nx.add(b)
Nx.tanh(output)

Visualizing a Linear Function: y = 2x

defmodule Mathy do
  def actual(x) do
    2 * x
  end
end

chart =
  Vl.new(width: 400, height: 400)
  |> Vl.mark(:line)
  |> Vl.encode_field(:x, "x", type: :quantitative)
  |> Vl.encode_field(:y, "y", type: :quantitative)
  |> Kino.VegaLite.new()
  |> Kino.render()

for i <- -5..5 do
  point = %{x: i, y: Mathy.actual(i)}
  Kino.VegaLite.push(chart, point)
  Process.sleep(25)
end

Solving linear function with Nx

defmodule Solver do
  import Nx.Defn

  @learning_rate 0.01
  @batch_size 64

  defn predict({w, b}, x) do
    w * x + b
  end

  defn mse(activations, x, y) do
    y_hat = predict(activations, x)

    (y - y_hat)
    |> Nx.pow(2)
    |> Nx.mean()
  end

  defn update({w, b} = activations, x, y) do
    {grad_w, grad_b} = grad(activations, &mse(&1, x, y))

    {
      w - grad_w * @learning_rate,
      b - grad_b * @learning_rate
    }
  end

  def train(data, epochs \\ 100) do
    Enum.reduce(1..epochs, random_parameters(), fn _i, acc ->
      data
      |> Enum.take(@batch_size)
      |> Enum.reduce(acc, fn batch, activations ->
        {x, y} = Enum.unzip(batch)
        x = Nx.tensor(x)
        y = Nx.tensor(y)
        update(activations, x, y)
      end)
    end)
  end

  def random_parameters do
    key = Nx.Random.key(12)
    {w, _new_key} = Nx.Random.normal(key, 0.0, 1.0)
    {b, _new_key} = Nx.Random.normal(key, 0.0, 1.0)

    {w, b}
  end

  def generate_data do
    num = :rand.uniform()
    {x, y} = {num, Mathy.actual(num)}
    {Nx.tensor(x), Nx.tensor(y)}
  end
end

Train in loop

data =
  Stream.repeatedly(fn ->
    Stream.map(0..64, fn _i ->
      num = :rand.uniform() * 10
      {num, Mathy.actual(num)}
    end)
  end)

model = Solver.train(data)

Solver.predict(model, 2)

Tanh graph

chart =
  Vl.new(width: 400, height: 400)
  |> Vl.mark(:line)
  |> Vl.encode_field(:x, "x", type: :quantitative)
  |> Vl.encode_field(:y, "y", type: :quantitative)
  |> Kino.VegaLite.new()
  |> Kino.render()

for i <- -10..10 do
  point = %{x: i, y: :math.tanh(i)}
  Kino.VegaLite.push(chart, point)
  Process.sleep(25)
end

Train NN for tanh with Axon

model =
  Axon.input("data")
  |> Axon.dense(4, activation: :relu)
  |> Axon.dense(1, activation: :tanh)

batch_size = 64

data =
  Stream.repeatedly(fn ->
    {x, _next_key} =
      12
      |> Nx.Random.key()
      |> Nx.Random.normal(-5.0, 5.0, shape: {64, 1})

    y = Nx.tanh(x)
    {x, y}
  end)

params =
  model
  |> Axon.Loop.trainer(:mean_squared_error, :sgd)
  |> Axon.Loop.run(data, %{}, epochs: 100, iterations: 100)

Axon.predict(model, params, Nx.tensor([[1]]))