Fix padding issue

zblanco · Jun 21, 2021 · b15be1f · b15be1f
1 parent d5d1e01
commit b15be1f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/lib/axon/layers.ex b/lib/axon/layers.ex
@@ -656,12 +656,21 @@ defmodule Axon.Layers do
         end
       )
 
+    padding =
+      transform({Nx.rank(input), opts[:padding]},
+        fn
+          {_, :same} -> :same
+          {_, :valid} -> :valid
+          {rank, padding} ->
+            List.duplicate({0, 0}, rank - 2) ++ padding
+        end)
+
     opts = transform(opts, &Keyword.delete(&1, :kernel_size))
 
     input
     |> Nx.window_max(window_dimensions,
       strides: strides,
-      padding: opts[:padding],
+      padding: padding,
       window_dilations: opts[:window_dilations]
     )
   end
@@ -720,12 +729,21 @@ defmodule Axon.Layers do
         end
       )
 
+    padding =
+      transform({Nx.rank(input), opts[:padding]},
+        fn
+          {_, :same} -> :same
+          {_, :valid} -> :valid
+          {rank, padding} ->
+            List.duplicate({0, 0}, rank - 2) ++ padding
+        end)
+
     opts = transform(opts, &Keyword.delete(&1, :kernel_size))
 
     input
     |> Nx.window_mean(window_dimensions,
       strides: strides,
-      padding: opts[:padding],
+      padding: padding,
       window_dilations: opts[:window_dilations]
     )
   end
@@ -804,6 +822,15 @@ defmodule Axon.Layers do
         end
       )
 
+    padding =
+      transform({Nx.rank(input), opts[:padding]},
+        fn
+          {_, :same} -> :same
+          {_, :valid} -> :valid
+          {rank, padding} ->
+            List.duplicate({0, 0}, rank - 2) ++ padding
+        end)
+
     norm = opts[:norm]
 
     opts =
@@ -815,7 +842,7 @@ defmodule Axon.Layers do
     |> Nx.power(norm)
     |> Nx.window_sum(window_dimensions,
       strides: strides,
-      padding: opts[:padding],
+      padding: padding,
       window_dilations: opts[:window_dilations]
     )
     |> Nx.power(Nx.divide(Nx.tensor(1, type: Nx.type(input)), norm))

diff --git a/lib/axon/training.ex b/lib/axon/training.ex
@@ -78,7 +78,7 @@ defmodule Axon.Training do
 
     step_fn = fn train_state, input, target ->
       {{preds, batch_loss}, gradients} =
-        Nx.Defn.Kernel.value_and_grad(
+        Nx.Defn.value_and_grad(
           train_state[:params],
           &objective_fn.(&1, input, target),
           fn x -> elem(x, 1) end