chore: use :polaris from Hex (#503)

examples/generative/mnist_gan.exs

@@ -132,8 +132,8 @@ defmodule MNISTGAN do
   end
 
   defp train_loop(d_model, g_model) do
-    {init_optim_d, optim_d} = Polaris.Optimizers.adam(2.0e-3, b1: 0.5)
-    {init_optim_g, optim_g} = Polaris.Optimizers.adam(2.0e-3, b1: 0.5)
+    {init_optim_d, optim_d} = Polaris.Optimizers.adam(learning_rate: 2.0e-3, b1: 0.5)
+    {init_optim_g, optim_g} = Polaris.Optimizers.adam(learning_rate: 2.0e-3, b1: 0.5)
 
     {init_d, d_model} = Axon.build(d_model, mode: :train)
     {init_g, g_model} = Axon.build(g_model, mode: :train)

examples/generative/text_generator.exs

@@ -95,7 +95,7 @@ defmodule TextGenerator do
 
     params =
       model
-      |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(0.001))
+      |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(learning_rate: 0.001))
       |> Axon.Loop.run(Stream.zip(train_data, train_labels), %{}, epochs: 20, compiler: EXLA)
 
     init_sequence = """

examples/structured/credit_card_fraud.exs

@@ -167,7 +167,7 @@ defmodule CreditCardFraud do
         reduction: :mean
       )
 
-    optimizer = Polaris.Optimizers.adam(1.0e-2)
+    optimizer = Polaris.Optimizers.adam(learning_rate: 1.0e-2)
 
     model
     |> train_model(loss, optimizer, batched_train)

examples/vision/horses_or_humans.exs

@@ -78,7 +78,7 @@ defmodule HorsesOrHumans do
 
   def run() do
     model = build_model({nil, 300, 300, 4}, [2, 0, 1]) |> IO.inspect()
-    optimizer = Polaris.Optimizers.adam(1.0e-4)
+    optimizer = Polaris.Optimizers.adam(learning_rate: 1.0e-4)
     centralized_optimizer = Polaris.Updates.compose(Polaris.Updates.centralize(), optimizer)
 
     data = data()

examples/vision/mnist.exs

@@ -37,7 +37,7 @@ defmodule Mnist do
 
   defp train_model(model, train_images, train_labels, epochs) do
     model
-    |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(0.005))
+    |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(learning_rate: 0.005))
     |> Axon.Loop.metric(:accuracy, "Accuracy")
     |> Axon.Loop.run(Stream.zip(train_images, train_labels), %{}, epochs: epochs, compiler: EXLA)
   end

guides/training_and_evaluation/custom_models_loss_optimizers.livemd

@@ -365,7 +365,7 @@ model =
   |> Axon.relu()
   |> Axon.dense(1)
 
-optimizer = {_init_optimizer_fn, _update_fn} = Polaris.Optimizers.sgd(1.0e-3)
+optimizer = {_init_optimizer_fn, _update_fn} = Polaris.Optimizers.sgd(learning_rate: 1.0e-3)
 
 model
 |> Axon.Loop.trainer(:mean_squared_error, optimizer)

lib/axon.ex

@@ -209,7 +209,7 @@ defmodule Axon do
 
       model_state =
         model
-        |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(0.005))
+        |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adamw(learning_rate: 0.005))
         |> Axon.Loop.run(train_data, epochs: 10, compiler: EXLA)
 
   See `Polaris.Updates` and `Axon.Loop` for a more in-depth treatment of
@@ -3032,7 +3032,7 @@ defmodule Axon do
         |> Axon.dense(1000, activation: :softmax)
 
       model
-      |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(0.005))
+      |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(learning_rate: 0.005))
       |> Axon.Loop.run(data, epochs: 10)
 
   When compiled, frozen parameters are wrapped in `Nx.Defn.Kernel.stop_grad/1`,
@@ -3104,7 +3104,7 @@ defmodule Axon do
         |> Axon.unfreeze(up: 25)
 
       model
-      |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(0.0005))
+      |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(learning_rate: 0.0005))
       |> Axon.Loop.run(data, epochs: 10)
 
   When compiled, frozen parameters are wrapped in `Nx.Defn.Kernel.stop_grad/1`,

lib/axon/loop.ex

@@ -654,15 +654,15 @@ defmodule Axon.Loop do
   ### Customizing Optimizer
 
       model
-      |> Axon.Loop.trainer(:binary_cross_entropy, Polaris.Optimizers.adam(0.05))
+      |> Axon.Loop.trainer(:binary_cross_entropy, Polaris.Optimizers.adam(learning_rate: 0.05))
       |> Axon.Loop.run(data)
 
   ### Custom loss
 
       loss_fn = fn y_true, y_pred -> Nx.cos(y_true, y_pred) end
 
       model
-      |> Axon.Loop.trainer(loss_fn, Polaris.Optimizers.rmsprop(0.01))
+      |> Axon.Loop.trainer(loss_fn, Polaris.Optimizers.rmsprop(learning_rate: 0.01))
       |> Axon.Loop.run(data)
 
   ### Multiple objectives with multi-output model

lib/axon/optimizers.ex

@@ -16,7 +16,7 @@ defmodule Axon.Optimizers do
 
     * [AdaBelief Optimizer: Adapting Stepsizes by the Belief in Observed Gradients](https://arxiv.org/abs/2010.07468)
   """
-  @deprecated "Use Polaris.Optimizers.adabelief/2 instead"
+  @deprecated "Use Polaris.Optimizers.adabelief/1 instead"
   def adabelief(learning_rate \\ 1.0e-3, opts \\ []) do
     Updates.scale_by_belief(opts)
     |> scale_by_learning_rate(learning_rate)
@@ -33,7 +33,7 @@ defmodule Axon.Optimizers do
 
     * [Adaptive Subgradient Methods for Online Learning and Stochastic Optimization](https://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf)
   """
-  @deprecated "Use Polaris.Optimizers.adagrad/2 instead"
+  @deprecated "Use Polaris.Optimizers.adagrad/1 instead"
   def adagrad(learning_rate \\ 1.0e-3, opts \\ []) do
     Updates.scale_by_rss(opts)
     |> scale_by_learning_rate(learning_rate)
@@ -53,7 +53,7 @@ defmodule Axon.Optimizers do
 
     * [Adam: A Method for Stochastic Optimization](https://arxiv.org/abs/1412.6980)
   """
-  @deprecated "Use Polaris.Optimizers.adam/2 instead"
+  @deprecated "Use Polaris.Optimizers.adam/1 instead"
   def adam(learning_rate \\ 1.0e-3, opts \\ []) do
     Updates.scale_by_adam(opts)
     |> scale_by_learning_rate(learning_rate)
@@ -70,7 +70,7 @@ defmodule Axon.Optimizers do
     * `:eps_root` - numerical stability term. Defaults to `0.0`
     * `:decay` - weight decay. Defaults to `0.0`
   """
-  @deprecated "Use Polaris.Optimizers.adamw/2 instead"
+  @deprecated "Use Polaris.Optimizers.adamw/1 instead"
   def adamw(learning_rate \\ 1.0e-3, opts \\ []) do
     {decay, opts} = Keyword.pop(opts, :decay, 0.0)
 
@@ -95,7 +95,7 @@ defmodule Axon.Optimizers do
 
     * [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://arxiv.org/abs/1904.00962)
   """
-  @deprecated "Use Polaris.Optimizers.lamb/2 instead"
+  @deprecated "Use Polaris.Optimizers.lamb/1 instead"
   def lamb(learning_rate \\ 1.0e-2, opts \\ []) do
     {decay, opts} = Keyword.pop(opts, :decay, 0.0)
     {min_norm, opts} = Keyword.pop(opts, :min_norm, 0.0)
@@ -114,7 +114,7 @@ defmodule Axon.Optimizers do
     * `:eta` - used to compute variance of noise distribution. Defaults to `0.1`
     * `:gamma` - used to compute variance of noise distribution. Defaults to `0.55`
   """
-  @deprecated "Use Polaris.Optimizers.noisy_sgd/2 instead"
+  @deprecated "Use Polaris.Optimizers.noisy_sgd/1 instead"
   def noisy_sgd(learning_rate \\ 1.0e-2, opts \\ []) do
     scale_by_learning_rate(learning_rate)
     |> Updates.add_noise(opts)
@@ -135,7 +135,7 @@ defmodule Axon.Optimizers do
 
     * [On the Variance of Adaptive Learning Rate and Beyond](https://arxiv.org/pdf/1908.03265.pdf)
   """
-  @deprecated "Use Polaris.Optimizers.radam/2 instead"
+  @deprecated "Use Polaris.Optimizers.radam/1 instead"
   def radam(learning_rate \\ 1.0e-3, opts \\ []) do
     Updates.scale_by_radam(opts)
     |> scale_by_learning_rate(learning_rate)
@@ -154,7 +154,7 @@ defmodule Axon.Optimizers do
     * `:decay` - EMA decay rate. Defaults to `0.9`
     * `:eps` - numerical stability term. Defaults to `1.0e-8`
   """
-  @deprecated "Use Polaris.Optimizers.rmsprop/2 instead"
+  @deprecated "Use Polaris.Optimizers.rmsprop/1 instead"
   def rmsprop(learning_rate \\ 1.0e-2, opts \\ []) do
     {centered, opts} = Keyword.pop(opts, :centered, false)
     {nesterov?, opts} = Keyword.pop(opts, :nesterov, false)
@@ -182,7 +182,7 @@ defmodule Axon.Optimizers do
       to value of this term.
     * `:nesterov` - whether or not to use nesterov momentum. Defaults to `false`
   """
-  @deprecated "Use Polaris.Optimizers.sgd/2 instead"
+  @deprecated "Use Polaris.Optimizers.sgd/1 instead"
   def sgd(learning_rate \\ 1.0e-2, opts \\ []) do
     momentum = opts[:momentum]
     nesterov? = opts[:nesterov] || false
@@ -210,7 +210,7 @@ defmodule Axon.Optimizers do
 
     * [Adaptive Methods for Nonconvex Optimization](https://papers.nips.cc/paper/2018/file/90365351ccc7437a1309dc64e4db32a3-Paper.pdf)
   """
-  @deprecated "Use Polaris.Optimizers.yogi/2 instead"
+  @deprecated "Use Polaris.Optimizers.yogi/1 instead"
   def yogi(learning_rate \\ 1.0e-2, opts \\ []) do
     Updates.scale_by_yogi(opts)
     |> scale_by_learning_rate(learning_rate)

mix.exs

@@ -42,7 +42,7 @@ defmodule Axon.MixProject do
       {:table_rex, "~> 3.1.1", optional: true},
       {:kino, "~> 0.7", optional: true},
       {:kino_vega_lite, "~> 0.1.7", optional: true},
-      {:polaris, "~> 0.1", github: "elixir-nx/polaris"}
+      {:polaris, "~> 0.1"}
     ]
   end
 

mix.lock

@@ -14,7 +14,7 @@
   "makeup_erlang": {:hex, :makeup_erlang, "0.1.1", "3fcb7f09eb9d98dc4d208f49cc955a34218fc41ff6b84df7c75b3e6e533cc65f", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "174d0809e98a4ef0b3309256cbf97101c6ec01c4ab0b23e926a9e17df2077cbb"},
   "nimble_parsec": {:hex, :nimble_parsec, "1.2.3", "244836e6e3f1200c7f30cb56733fd808744eca61fd182f731eac4af635cc6d0b", [:mix], [], "hexpm", "c8d789e39b9131acf7b99291e93dae60ab48ef14a7ee9d58c6964f59efb570b0"},
   "nx": {:hex, :nx, "0.5.1", "118134b8c97c2a8f86c87aa8434994c1cbbe139a306b89cca04e08dd46228067", [:mix], [{:complex, "~> 0.5", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "ceb8fbbe19b3c4252a7188d8b0e059fac9da0f4a4f3bb770fc665fdd0b29f0c5"},
-  "polaris": {:git, "https://github.com/elixir-nx/polaris.git", "d306b4084159b50d07293d3cfde9f2a1a2ef0d7b", []},
+  "polaris": {:hex, :polaris, "0.1.0", "dca61b18e3e801ecdae6ac9f0eca5f19792b44a5cb4b8d63db50fc40fc038d22", [:mix], [{:nx, "~> 0.5", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "13ef2b166650e533cb24b10e2f3b8ab4f2f449ba4d63156e8c569527f206e2c2"},
   "table": {:hex, :table, "0.1.2", "87ad1125f5b70c5dea0307aa633194083eb5182ec537efc94e96af08937e14a8", [:mix], [], "hexpm", "7e99bc7efef806315c7e65640724bf165c3061cdc5d854060f74468367065029"},
   "table_rex": {:hex, :table_rex, "3.1.1", "0c67164d1714b5e806d5067c1e96ff098ba7ae79413cc075973e17c38a587caa", [:mix], [], "hexpm", "678a23aba4d670419c23c17790f9dcd635a4a89022040df7d5d772cb21012490"},
   "telemetry": {:hex, :telemetry, "1.2.1", "68fdfe8d8f05a8428483a97d7aab2f268aaff24b49e0f599faa091f1d4e7f61c", [:rebar3], [], "hexpm", "dad9ce9d8effc621708f99eac538ef1cbe05d6a874dd741de2e689c47feafed5"},

notebooks/generative/fashionmnist_vae.livemd

@@ -251,7 +251,7 @@ end
 
 params =
   model
-  |> Axon.Loop.trainer(:mean_squared_error, Polaris.Optimizers.adamw(0.001))
+  |> Axon.Loop.trainer(:mean_squared_error, Polaris.Optimizers.adamw(learning_rate: 0.001))
   |> KinoAxon.kino_early_stop()
   |> Axon.Loop.handle(:iteration_completed, render_example_handler, every: 450)
   |> Axon.Loop.validate(model, test_data)
@@ -474,7 +474,7 @@ end
 
 params =
   model
-  |> Axon.Loop.trainer(&CustomLoss.loss/2, Polaris.Optimizers.adam(0.001))
+  |> Axon.Loop.trainer(&CustomLoss.loss/2, Polaris.Optimizers.adam(learning_rate: 0.001))
   |> KinoAxon.kino_early_stop()
   |> Axon.Loop.handle(:epoch_completed, render_example_handler)
   |> Axon.Loop.validate(model, test_data)

notebooks/generative/mnist_autoencoder_using_kino.livemd

@@ -197,7 +197,7 @@ Looks right (and tricky). Let's see how the model does.
 ```elixir
 params =
   model
-  |> Axon.Loop.trainer(:mean_squared_error, Polaris.Optimizers.adamw(0.001))
+  |> Axon.Loop.trainer(:mean_squared_error, Polaris.Optimizers.adamw(learning_rate: 0.001))
   |> Axon.Loop.validate(model, test_data)
   |> Axon.Loop.run(train_data, %{}, epochs: 20, compiler: EXLA)
 
@@ -312,7 +312,7 @@ end
 
 params =
   model
-  |> Axon.Loop.trainer(:mean_squared_error, Polaris.Optimizers.adamw(0.001))
+  |> Axon.Loop.trainer(:mean_squared_error, Polaris.Optimizers.adamw(learning_rate: 0.001))
   |> Axon.Loop.handle(:iteration_completed, render_example_handler, every: 450)
   |> Axon.Loop.validate(model, test_data)
   |> Axon.Loop.run(train_data, %{}, epochs: 20, compiler: EXLA)

notebooks/structured/credit_card_fraud.livemd

@@ -151,7 +151,7 @@ loss =
     reduction: :mean
   )
 
-optimizer = Polaris.Optimizers.adam(1.0e-2)
+optimizer = Polaris.Optimizers.adam(learning_rate: 1.0e-2)
 
 params =
   model

notebooks/text/lstm_generation.livemd

@@ -171,7 +171,7 @@ IO.puts("Total batches: #{Enum.count(train_batches)}")
 
 params =
   model
-  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(0.001))
+  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(learning_rate: 0.001))
   |> Axon.Loop.run(Stream.zip(train_batches, result_batches), %{}, epochs: 20, compiler: EXLA)
 
 :ok
@@ -250,7 +250,7 @@ IO.puts("Total batches: #{Enum.count(train_batches)}")
 
 new_params =
   new_model
-  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(0.001))
+  |> Axon.Loop.trainer(:categorical_cross_entropy, Polaris.Optimizers.adam(learning_rate: 0.001))
   |> Axon.Loop.run(Stream.zip(train_batches, result_batches), %{}, epochs: 50, compiler: EXLA)
 
 :ok

notebooks/vision/horses_or_humans.livemd

@@ -199,7 +199,7 @@ It's time to train our model. We specify the loss, optimizer and choose accuracy
 ```elixir
 data = HorsesHumans.DataProcessing.data_stream(files, batch_size)
 
-optimizer = Polaris.Optimizers.adam(1.0e-4)
+optimizer = Polaris.Optimizers.adam(learning_rate: 1.0e-4)
 
 params =
   model