CompileError when using combination of Axon.dropout and some optimizers with Axon v0.4.0

[nickgnd]

Hey 👋
as always, thanks for the great work you are doing, so inspiring! 🙇

I'm reaching out because I'm getting an error with the latest version of axon (v0.4.0) when using Axon.dropout in the model and then one of the following optimizers in the Axon.trainer api:

• rmsprop
• adagrad
• yogi

While, it works fine with the previous version of Axon (v0.3.1).

It might be the same issue described in #443 ?

---

Here the snippet that you can use to reproduce the issue:

Mix.install(
  [
    {:exla, "~> 0.4"},
    {:nx, "~> 0.4"},
    {:axon, "~> 0.4"}
    # {:axon, "~> 0.3.0"} # to install the prev. version of Axon
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

# Generate the data

key = Nx.Random.key(42)

{inputs, _new_key} = Nx.Random.normal(key, 0, 1, shape: {1000, 2}, type: :f32)

labels =
  Enum.map(0..999, fn _ -> Enum.random([0, 1]) end)
  |> Nx.tensor()
  |> Nx.reshape({:auto, 1})
  |> Nx.equal(Nx.tensor([0, 1]))

# Create model and train

model =
  Axon.input("data", shape: {nil, 2})
  |> Axon.dense(100, activation: :sigmoid)
  |> Axon.dropout()
  |> Axon.dense(2, activation: :softmax)

# Erros with `rmsprop`, `adagrad` and `yogi` optimizers
model
|> Axon.Loop.trainer(:categorical_cross_entropy, :rmsprop)
|> Axon.Loop.run([{inputs, labels}], %{}, epochs: 10, compiler: EXLA)

When running the livebook I'm receiving this error:

** (CompileError) /Users/nicolo.gnudi/Library/Caches/mix/installs/elixir-1.14.2-erts-13.1.2/7aad6e1955530b5d7550df33b2b3f34b/deps/axon/lib/axon/loop.ex:462: the do-block in while must return tensors with the same shape, type, and names as the initial arguments.

Body matches template:

{%{"dense_0" => %{"bias" => #Nx.Tensor<
       f32[100]
     >, "kernel" => #Nx.Tensor<
       f32[2][100]
     >}, "dense_1" => %{"bias" => #Nx.Tensor<
       f32[2]
     >, "kernel" => #Nx.Tensor<
       f32[100][2]
     >}, "dropout_0" => %{"key" => #Nx.Tensor<
       f32[2]
     >}}, %{"dense_0" => %{"bias" => #Nx.Tensor<
       f32[100]
     >, "kernel" => #Nx.Tensor<
       f32[2][100]
     >}, "dense_1" => %{"bias" => #Nx.Tensor<
       f32[2]
     >, "kernel" => #Nx.Tensor<
       f32[100][2]
     >}, "dropout_0" => %{"key" => #Nx.Tensor<
       u32[2]
     >}}, %{"dropout_0" => %{"key" => #Nx.Tensor<
       u32[2]
     >}}, {%{scale: #Nx.Tensor<
      f32
    >}, %{nu: %{"dense_0" => %{"bias" => #Nx.Tensor<
          f32[100]
        >, "kernel" => #Nx.Tensor<
          f32[2][100]
        >}, "dense_1" => %{"bias" => #Nx.Tensor<
          f32[2]
        >, "kernel" => #Nx.Tensor<
          f32[100][2]
        >}, "dropout_0" => %{"key" => #Nx.Tensor<
          f32[2]
        >}}}}, %{"dense_0" => %{"bias" => #Nx.Tensor<
       f32[100]
     >, "kernel" => #Nx.Tensor<
       f32[2][100]
     >}, "dense_1" => %{"bias" => #Nx.Tensor<
       f32[2]
     >, "kernel" => #Nx.Tensor<
       f32[100][2]
     >}, "dropout_0" => %{"key" => #Nx.Tensor<
       f32[2]
     >}}, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64
 >}

and initial argument has template:

{%{"dense_0" => %{"bias" => #Nx.Tensor<
       f32[100]
     >, "kernel" => #Nx.Tensor<
       f32[2][100]
     >}, "dense_1" => %{"bias" => #Nx.Tensor<
       f32[2]
     >, "kernel" => #Nx.Tensor<
       f32[100][2]
     >}, "dropout_0" => %{"key" => #Nx.Tensor<
       f32[2]
     >}}, %{"dense_0" => %{"bias" => #Nx.Tensor<
       f32[100]
     >, "kernel" => #Nx.Tensor<
       f32[2][100]
     >}, "dense_1" => %{"bias" => #Nx.Tensor<
       f32[2]
     >, "kernel" => #Nx.Tensor<
       f32[100][2]
     >}, "dropout_0" => %{"key" => #Nx.Tensor<
       u32[2]
     >}}, %{"dropout_0" => %{"key" => #Nx.Tensor<
       u32[2]
     >}}, {%{scale: #Nx.Tensor<
      f32
    >}, %{nu: %{"dense_0" => %{"bias" => #Nx.Tensor<
          f32[100]
        >, "kernel" => #Nx.Tensor<
          f32[2][100]
        >}, "dense_1" => %{"bias" => #Nx.Tensor<
          f32[2]
        >, "kernel" => #Nx.Tensor<
          f32[100][2]
        >}, "dropout_0" => %{"key" => #Nx.Tensor<
          u32[2]
        >}}}}, %{"dense_0" => %{"bias" => #Nx.Tensor<
       f32[100]
     >, "kernel" => #Nx.Tensor<
       f32[2][100]
     >}, "dense_1" => %{"bias" => #Nx.Tensor<
       f32[2]
     >, "kernel" => #Nx.Tensor<
       f32[100][2]
     >}, "dropout_0" => %{"key" => #Nx.Tensor<
       f32[2]
     >}}, #Nx.Tensor<
   s64
 >, #Nx.Tensor<
   s64
 >}

    (nx 0.4.2) lib/nx/defn/expr.ex:483: Nx.Defn.Expr.compatible_while!/4
    (nx 0.4.2) lib/nx/defn/expr.ex:354: Nx.Defn.Expr.defn_while/6
    (axon 0.4.0) lib/axon/loop.ex:462: Axon.Loop."__defn:accumulate_gradients__"/8
    (axon 0.4.0) lib/axon/loop.ex:414: anonymous fn/6 in Axon.Loop.train_step/4
    (axon 0.4.0) lib/axon/loop.ex:1921: anonymous fn/6 in Axon.Loop.build_batch_fn/2
    (nx 0.4.2) lib/nx/defn/compiler.ex:142: Nx.Defn.Compiler.runtime_fun/3
    (exla 0.4.2) lib/exla/defn.ex:378: anonymous fn/4 in EXLA.Defn.compile/8
    /Users/nicolo.gnudi/code/personal/programming_machine_learning/17_overfitting/debugging_dropout.livemd#cell:msdfdrfgdvu22mzha37twp5fiwyqg677:24: (file)

For what I can see in the printed error, there is a type mismatch in the dropout layer, in the 2nd element of the tuple (I mean %{nu: …}):

In "Body matches template" there is:

>}, "dropout_0" => %{"key" => #Nx.Tensor<
  f32[2]
>}}}},

while in "initial argument has template" there is:

>}, "dropout_0" => %{"key" => #Nx.Tensor<
  u32[2]
>}}}}, 

---

Let me know if there is anything else I can do to ease the debugging 🙇
Cheers ✌️

My setup:

• macOS Monterey 12.6
• erlang 25.1.2
• elixir 1.14.2-otp-25
• livebook 0.8.1
• processor: 2.3 GHz Quad-Core Intel Core i7
• memory: 32 GB
• graphic card: Intel Iris Plus Graphics 1536 MB

[seanmor5]

I just released v0.4.1 with a fix :)

[nickgnd]

@seanmor5 wow, so quick! Thank youuuuu 🙏 💜