docs: adds linear regression example

.github/workflows/ci.yml

@@ -34,3 +34,4 @@ jobs:
       - run: mix compile --warnings-as-errors
         if: ${{ matrix.lint }}
       - run: mix test
+      - run: mix run examples/linear_regression.exs

examples/linear_regression.exs

@@ -0,0 +1,75 @@
+defmodule LinearRegression do
+  import Nx.Defn
+
+  # y = mx + b
+  defn init_random_params do
+    key = Nx.Random.key(42)
+    {m, new_key} = Nx.Random.normal(key, 0.0, 0.1, shape: {1, 1})
+    {b, _new_key} = Nx.Random.normal(new_key, 0.0, 0.1, shape: {1})
+    {m, b}
+  end
+
+  defn predict({m, b}, inp) do
+    Nx.dot(inp, m) + b
+  end
+
+  # MSE Loss
+  defn loss({m, b}, inp, tar) do
+    Nx.mean(Nx.pow(tar - predict({m, b}, inp), 2))
+  end
+
+  defn update({m, b} = params, inp, tar, step) do
+    {grad_m, grad_b} = grad(params, &loss(&1, inp, tar))
+    {m - grad_m * step, b - grad_b * step}
+  end
+
+  def train(params, epochs, lin_fn) do
+    data =
+      Stream.repeatedly(fn -> for _ <- 1..32, do: :rand.uniform() * 10 end)
+      |> Stream.map(fn x -> Enum.zip(x, Enum.map(x, lin_fn)) end)
+
+    for _ <- 1..epochs, reduce: params do
+      acc ->
+        data
+        |> Enum.take(200)
+        |> Enum.reduce(
+          acc,
+          fn batch, cur_params ->
+            {inp, tar} = Enum.unzip(batch)
+            x = Nx.reshape(Nx.tensor(inp), {32, 1})
+            y = Nx.reshape(Nx.tensor(tar), {32, 1})
+            update(cur_params, x, y, 0.001)
+          end
+        )
+    end
+  end
+end
+
+Nx.default_backend(Candlex.Backend)
+
+params = LinearRegression.init_random_params()
+m = :rand.normal(0.0, 10.0)
+b = :rand.normal(0.0, 5.0)
+IO.puts("Target m: #{m} Target b: #{b}\n")
+
+lin_fn = fn x -> m * x + b end
+epochs = 100
+
+# These will be very close to the above coefficients
+{time, {trained_m, trained_b}} = :timer.tc(LinearRegression, :train, [params, epochs, lin_fn])
+
+trained_m =
+  trained_m
+  |> Nx.squeeze()
+  |> Nx.backend_transfer()
+  |> Nx.to_number()
+
+trained_b =
+  trained_b
+  |> Nx.squeeze()
+  |> Nx.backend_transfer()
+  |> Nx.to_number()
+
+IO.puts("Trained in #{time / 1_000_000} sec.")
+IO.puts("Trained m: #{trained_m} Trained b: #{trained_b}\n")
+IO.puts("Accuracy m: #{m - trained_m} Accuracy b: #{b - trained_b}")