Fix mismatch for certain optimizers with dropout (#458)

Author: seanmor5

lib/axon/shared.ex

@@ -135,11 +135,13 @@ defmodule Axon.Shared do
   @doc """
   Creates a fulls-like tuple of inputs.
   """
-  deftransform fulls_like(params, value) do
+  deftransform fulls_like(params, value, opts \\ []) do
+    opts = Keyword.validate!(opts, [:type])
     fun = Axon.Initializers.full(value)
 
     deep_new(params, fn x ->
-      fun.(Nx.shape(x), Nx.type(x))
+      type = opts[:type] || Nx.type(x)
+      fun.(Nx.shape(x), type)
     end)
   end
 

lib/axon/updates.ex

@@ -219,7 +219,7 @@ defmodule Axon.Updates do
   end
 
   defnp init_scale_by_rss(params, value) do
-    sum_of_squares = fulls_like(params, value)
+    sum_of_squares = fulls_like(params, value, type: :f32)
     %{sum_of_squares: sum_of_squares}
   end
 
@@ -278,7 +278,7 @@ defmodule Axon.Updates do
   end
 
   defnp init_scale_by_rms(params, scale) do
-    nu = fulls_like(params, scale)
+    nu = fulls_like(params, scale, type: :f32)
     %{nu: nu}
   end
 
@@ -395,7 +395,7 @@ defmodule Axon.Updates do
 
   defnp init_scale_by_stddev(params, value) do
     mu = zeros_like(params, type: :f32)
-    nu = fulls_like(params, value)
+    nu = fulls_like(params, value, type: :f32)
     %{mu: mu, nu: nu}
   end
 
@@ -860,7 +860,7 @@ defmodule Axon.Updates do
   end
 
   defnp init_scale_by_yogi(params, value) do
-    value = fulls_like(params, value)
+    value = fulls_like(params, value, type: :f32)
     mu = value
     nu = value
     count = Nx.tensor(0)

mix.exs

@@ -2,7 +2,7 @@ defmodule Axon.MixProject do
   use Mix.Project
 
   @source_url "https://github.com/elixir-nx/axon"
-  @version "0.4.0"
+  @version "0.4.1"
 
   def project do
     [

test/axon/integration_test.exs

@@ -184,4 +184,58 @@ defmodule Axon.IntegrationTest do
       assert_equal(step_state1, step_state2)
     end)
   end
+
+  test "dropout with certain optimizers regression test" do
+    {train, _test} = get_test_data(100, 0, 10, {10}, 2, 1337)
+
+    train =
+      train
+      |> Stream.map(fn {xs, ys} ->
+        {xs, one_hot(ys, num_classes: 2)}
+      end)
+      |> Enum.to_list()
+
+    [{x_test, _}] = Enum.take(train, 1)
+
+    model =
+      Axon.input("input")
+      |> Axon.dense(16)
+      |> Axon.dropout(rate: 0.1)
+      |> Axon.dense(2, activation: :softmax)
+
+    optimizers = [
+      Axon.Optimizers.rmsprop(5.0e-3, centered: true),
+      Axon.Optimizers.rmsprop(5.0e-3, centered: false),
+      :adagrad,
+      :yogi
+    ]
+
+    ExUnit.CaptureIO.capture_io(fn ->
+      for optim <- optimizers do
+        results =
+          model
+          |> Axon.Loop.trainer(:categorical_cross_entropy, optim)
+          # TODO: Fix default output transform
+          |> Map.update(:output_transform, nil, fn _ -> & &1 end)
+          |> Axon.Loop.metric(:accuracy)
+          |> Axon.Loop.validate(model, train)
+          |> Axon.Loop.run(train, %{}, epochs: 10)
+
+        assert %{step_state: %{model_state: model_state}, metrics: %{9 => last_epoch_metrics}} =
+                 results
+
+        eval_results =
+          model
+          |> Axon.Loop.evaluator()
+          |> Axon.Loop.metric(:accuracy)
+          |> Axon.Loop.run(train, model_state)
+
+        assert %{0 => %{"accuracy" => final_model_val_accuracy}} = eval_results
+
+        assert_greater_equal(last_epoch_metrics["validation_accuracy"], 0.7)
+        assert_all_close(final_model_val_accuracy, last_epoch_metrics["validation_accuracy"])
+        assert Nx.shape(Axon.predict(model, model_state, x_test)) == {10, 2}
+      end
+    end)
+  end
 end