Fix deprecation warnings with new Nx (#454)

Author: seanmor5

lib/axon/activations.ex

@@ -525,7 +525,8 @@ defmodule Axon.Activations do
   defn relu(x) do
     custom_grad(
       Nx.max(x, 0),
-      fn _ans, g -> [{x, Nx.select(Nx.greater(x, 0), g, Nx.broadcast(0, g))}] end
+      [x],
+      fn g -> [{x, Nx.select(Nx.greater(x, 0), g, Nx.broadcast(0, g))}] end
     )
   end
 

lib/axon/compiler.ex

@@ -889,7 +889,7 @@ defmodule Axon.Compiler do
 
       if event? and mode? do
         if on_event == :backward do
-          Nx.Defn.Kernel.custom_grad(expr, fn _ans, g ->
+          Nx.Defn.Kernel.custom_grad(expr, [expr], fn g ->
             hooked_g = Nx.Defn.Kernel.hook(g, hook_fn)
             [{expr, hooked_g}]
           end)

lib/axon/initializers.ex

@@ -174,20 +174,16 @@ defmodule Axon.Initializers do
 
   """
   def uniform(opts \\ []) do
+    opts = Keyword.validate!(opts, scale: 1.0e-2)
+    scale = Keyword.fetch!(opts, :scale)
+
     fn shape, type, key ->
-      scale = opts[:scale] || 1.0e-2
-      uniform_impl(key, shape: shape, type: type, scale: scale)
+      uniform_impl(key, scale, shape: shape, type: type)
     end
   end
 
-  defnp uniform_impl(key, opts \\ []) do
-    opts = keyword!(opts, [:shape, type: {:f, 32}, scale: 1.0e-2])
-    shape = Nx.shape(opts[:shape])
-
-    Nx.Random.uniform_split(key, Nx.negate(opts[:scale]), opts[:scale],
-      type: opts[:type],
-      shape: shape
-    )
+  defnp uniform_impl(key, scale, opts) do
+    Nx.Random.uniform_split(key, Nx.negate(scale), scale, opts)
   end
 
   @doc """
@@ -216,18 +212,15 @@ defmodule Axon.Initializers do
 
   """
   def normal(opts \\ []) do
+    opts = Keyword.validate!(opts, scale: 1.0e-2, mean: 0.0)
+    scale = Keyword.fetch!(opts, :scale)
+    mean = Keyword.fetch!(opts, :mean)
+
     fn shape, type, key ->
-      scale = opts[:scale] || 1.0e-2
-      mean = opts[:mean] || 0.0
-      normal_impl(key, shape: shape, type: type, scale: scale, mean: mean)
+      Nx.Random.normal_split(key, mean, scale, type: type, shape: shape)
     end
   end
 
-  defnp normal_impl(key, opts \\ []) do
-    opts = keyword!(opts, [:shape, type: {:f, 32}, scale: 1.0e-2, mean: 0.0])
-    Nx.Random.normal_split(key, opts[:mean], opts[:scale], shape: opts[:shape], type: opts[:type])
-  end
-
   @doc """
   Initializes parameters with the Lecun uniform initializer.
 
@@ -261,25 +254,21 @@ defmodule Axon.Initializers do
 
   """
   def lecun_uniform(opts \\ []) do
+    opts = Keyword.validate!(opts, scale: 1.0)
+    scale = Keyword.fetch!(opts, :scale)
+
     fn shape, type, key ->
-      scale = opts[:scale] || 1.0
-      lecun_uniform_impl(key, shape: shape, type: type, scale: scale)
+      variance_scaling_impl(
+        key,
+        scale,
+        shape: shape,
+        type: type,
+        mode: :fan_in,
+        distribution: :uniform
+      )
     end
   end
 
-  defnp lecun_uniform_impl(key, opts \\ []) do
-    opts = keyword!(opts, [:shape, type: {:f, 32}, scale: 1.0])
-
-    variance_scaling_impl(
-      key,
-      shape: opts[:shape],
-      type: opts[:type],
-      scale: opts[:scale],
-      mode: :fan_in,
-      distribution: :uniform
-    )
-  end
-
   @doc """
   Initializes parameters with the Lecun normal initializer.
 
@@ -313,25 +302,21 @@ defmodule Axon.Initializers do
 
   """
   def lecun_normal(opts \\ []) do
+    opts = Keyword.validate!(opts, scale: 1.0)
+    scale = Keyword.fetch!(opts, :scale)
+
     fn shape, type, key ->
-      scale = opts[:scale] || 1.0
-      lecun_normal_impl(key, shape: shape, type: type, scale: scale)
+      variance_scaling_impl(
+        key,
+        scale,
+        shape: shape,
+        type: type,
+        mode: :fan_in,
+        distribution: :truncated_normal
+      )
     end
   end
 
-  defnp lecun_normal_impl(key, opts \\ []) do
-    opts = keyword!(opts, [:shape, type: {:f, 32}, scale: 1.0])
-
-    variance_scaling_impl(
-      key,
-      shape: opts[:shape],
-      type: opts[:type],
-      scale: opts[:scale],
-      mode: :fan_in,
-      distribution: :truncated_normal
-    )
-  end
-
   @doc """
   Initializes parameters with the Glorot uniform initializer.
 
@@ -368,25 +353,21 @@ defmodule Axon.Initializers do
 
   """
   def glorot_uniform(opts \\ []) do
+    opts = Keyword.validate!(opts, scale: 1.0)
+    scale = Keyword.fetch!(opts, :scale)
+
     fn shape, type, key ->
-      scale = opts[:scale] || 1.0
-      glorot_uniform_impl(key, shape: shape, type: type, scale: scale)
+      variance_scaling_impl(
+        key,
+        scale,
+        shape: shape,
+        type: type,
+        mode: :fan_avg,
+        distribution: :uniform
+      )
     end
   end
 
-  defnp glorot_uniform_impl(key, opts \\ []) do
-    opts = keyword!(opts, [:shape, type: {:f, 32}, scale: 1.0])
-
-    variance_scaling_impl(
-      key,
-      shape: opts[:shape],
-      type: opts[:type],
-      scale: opts[:scale],
-      mode: :fan_avg,
-      distribution: :uniform
-    )
-  end
-
   @doc """
   Initializes parameters with the Glorot normal initializer.
 
@@ -423,25 +404,21 @@ defmodule Axon.Initializers do
 
   """
   def glorot_normal(opts \\ []) do
+    opts = Keyword.validate!(opts, scale: 1.0)
+    scale = Keyword.fetch!(opts, :scale)
+
     fn shape, type, key ->
-      scale = opts[:scale] || 1.0
-      glorot_normal_impl(key, shape: shape, type: type, scale: scale)
+      variance_scaling_impl(
+        key,
+        scale,
+        shape: shape,
+        type: type,
+        mode: :fan_avg,
+        distribution: :truncated_normal
+      )
     end
   end
 
-  defnp glorot_normal_impl(key, opts \\ []) do
-    opts = keyword!(opts, [:shape, type: {:f, 32}, scale: 1.0])
-
-    variance_scaling_impl(
-      key,
-      shape: opts[:shape],
-      type: opts[:type],
-      scale: opts[:scale],
-      mode: :fan_avg,
-      distribution: :truncated_normal
-    )
-  end
-
   @doc """
   Initializes parameters with the He uniform initializer.
 
@@ -475,25 +452,21 @@ defmodule Axon.Initializers do
 
   """
   def he_uniform(opts \\ []) do
+    opts = Keyword.validate!(opts, scale: 2.0)
+    scale = Keyword.fetch!(opts, :scale)
+
     fn shape, type, key ->
-      scale = opts[:scale] || 2.0
-      he_uniform_impl(key, shape: shape, type: type, scale: scale)
+      variance_scaling_impl(
+        key,
+        scale,
+        shape: shape,
+        type: type,
+        mode: :fan_in,
+        distribution: :uniform
+      )
     end
   end
 
-  defnp he_uniform_impl(key, opts \\ []) do
-    opts = keyword!(opts, [:shape, type: {:f, 32}, scale: 2.0])
-
-    variance_scaling_impl(
-      key,
-      shape: opts[:shape],
-      type: opts[:type],
-      scale: opts[:scale],
-      mode: :fan_in,
-      distribution: :uniform
-    )
-  end
-
   @doc """
   Initializes parameters with the He normal initializer.
 
@@ -527,25 +500,21 @@ defmodule Axon.Initializers do
 
   """
   def he_normal(opts \\ []) do
+    opts = Keyword.validate!(opts, scale: 2.0)
+    scale = Keyword.fetch!(opts, :scale)
+
     fn shape, type, key ->
-      scale = opts[:scale] || 2.0
-      he_normal_impl(key, shape: shape, type: type, scale: scale)
+      variance_scaling_impl(
+        key,
+        scale,
+        shape: shape,
+        type: type,
+        mode: :fan_in,
+        distribution: :truncated_normal
+      )
     end
   end
 
-  defnp he_normal_impl(key, opts \\ []) do
-    opts = keyword!(opts, [:shape, type: {:f, 32}, scale: 2.0])
-
-    variance_scaling_impl(
-      key,
-      shape: opts[:shape],
-      type: opts[:type],
-      scale: opts[:scale],
-      mode: :fan_in,
-      distribution: :truncated_normal
-    )
-  end
-
   @doc """
   Initializes parameters with variance scaling according to
   the given distribution and mode.
@@ -586,30 +555,29 @@ defmodule Axon.Initializers do
 
   """
   def variance_scaling(opts \\ []) do
-    fn shape, type, key ->
-      scale = opts[:scale] || 1.0
-      mode = opts[:mode] || :fan_in
-      distribution = opts[:distribution] || :normal
+    opts = Keyword.validate!(opts, scale: 1.0, mode: :fan_in, distribution: :normal)
+    scale = Keyword.fetch!(opts, :scale)
+    mode = Keyword.fetch!(opts, :mode)
+    distribution = Keyword.fetch!(opts, :distribution)
 
+    fn shape, type, key ->
       variance_scaling_impl(
         key,
+        scale,
         shape: shape,
         type: type,
-        scale: scale,
         mode: mode,
         distribution: distribution
       )
     end
   end
 
-  defnp variance_scaling_impl(key, opts \\ []) do
-    opts =
-      keyword!(opts, [:shape, type: {:f, 32}, scale: 1.0, mode: :fan_in, distribution: :normal])
+  defnp variance_scaling_impl(key, scale, opts \\ []) do
+    opts = keyword!(opts, [:shape, type: {:f, 32}, mode: :fan_in, distribution: :normal])
 
     fans = compute_fans(opts[:shape])
     denominator = compute_denominator(fans, opts[:mode])
-
-    variance = Nx.divide(Nx.tensor(opts[:scale], type: opts[:type]), Nx.max(denominator, 1.0))
+    variance = Nx.as_type(scale, opts[:type]) / Nx.max(denominator, 1.0)
 
     apply_distribution(key, opts[:distribution], variance, shape: opts[:shape], type: opts[:type])
   end

lib/axon/layers.ex

@@ -1609,7 +1609,7 @@ defmodule Axon.Layers do
   @doc type: :dropout
   defn dropout(input, key, opts \\ []) do
     opts = keyword!(opts, [:rate, noise_shape: Nx.shape(input), mode: :inference])
-    keep_prob = Nx.tensor(1, type: Nx.type(input)) - Nx.tensor(opts[:rate], type: Nx.type(input))
+    keep_prob = Nx.tensor(1, type: Nx.type(input)) - Nx.as_type(opts[:rate], Nx.type(input))
 
     {rand, new_key} =
       Nx.Random.uniform(key, 0, 1, shape: opts[:noise_shape], type: Nx.type(input))

lib/axon/schedules.ex

@@ -22,7 +22,6 @@ defmodule Axon.Schedules do
   """
 
   import Nx.Defn
-  import Axon.Shared
 
   @doc """
   Linear decay schedule.
@@ -86,7 +85,7 @@ defmodule Axon.Schedules do
 
     init_value = opts[:init_value]
     rate = opts[:decay_rate]
-    staircase? = to_predicate(opts[:staircase])
+    staircase? = opts[:staircase]
     k = opts[:transition_steps]
     start = opts[:transition_begin]
 
@@ -166,7 +165,7 @@ defmodule Axon.Schedules do
 
   defnp apply_constant(_step, opts \\ []) do
     opts = keyword!(opts, init_value: 0.01)
-    Nx.tensor(opts[:init_value])
+    opts[:init_value]
   end
 
   @doc ~S"""