Annotate activation functions with Nx blocks

.tool-versions

@@ -0,0 +1,2 @@
+erlang 27.3
+elixir 1.20.2-otp-27

lib/axon/activations.ex

@@ -44,6 +44,7 @@ defmodule Axon.Activations do
   """
 
   import Nx.Defn
+  import Axon.Block
   import Axon.Shared
 
   @doc ~S"""
@@ -83,7 +84,7 @@ defmodule Axon.Activations do
     * [Continuously Differentiable Exponential Linear Units](https://arxiv.org/pdf/1704.07483.pdf)
 
   """
-  defn celu(x, opts \\ []) do
+  defblock CELU, celu(x, opts \\ []) do
     opts = keyword!(opts, alpha: 1.0)
     validate_celu_alpha!(opts[:alpha])
 
@@ -128,7 +129,7 @@ defmodule Axon.Activations do
     * [Fast and Accurate Deep Network Learning by Exponential Linear Units (ELUs)](https://arxiv.org/abs/1511.07289)
 
   """
-  defn elu(x, opts \\ []) do
+  defblock ELU, elu(x, opts \\ []) do
     opts = keyword!(opts, alpha: 1.0)
     x_hat = Nx.select(Nx.greater(x, 0), 0, x)
     Nx.select(Nx.greater(x, 0), x, opts[:alpha] * Nx.expm1(x_hat))
@@ -157,7 +158,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn exp(x) do
+  defblock exp(x) do
     Nx.exp(x)
   end
 
@@ -188,7 +189,7 @@ defmodule Axon.Activations do
     * [Gaussian Error Linear Units (GELUs)](https://arxiv.org/abs/1606.08415)
 
   """
-  defn gelu(x) do
+  defblock GeLU, gelu(x) do
     sqrt2 = Nx.sqrt(Nx.tensor(2, type: Nx.type(x)))
 
     x
@@ -220,7 +221,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn hard_sigmoid(x, opts \\ []) do
+  defblock hard_sigmoid(x, opts \\ []) do
     opts = keyword!(opts, alpha: 0.2, beta: 0.2)
 
     x
@@ -255,7 +256,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn hard_silu(x, opts \\ []) do
+  defblock HardSiLU, hard_silu(x, opts \\ []) do
     x
     |> hard_sigmoid(opts)
     |> Nx.multiply(x)
@@ -284,7 +285,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn hard_tanh(x) do
+  defblock hard_tanh(x) do
     Nx.select(
       Nx.greater(x, 1),
       1,
@@ -319,7 +320,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn leaky_relu(x, opts \\ []) do
+  defblock LeakyReLU, leaky_relu(x, opts \\ []) do
     opts = keyword!(opts, alpha: 1.0e-2)
     Nx.select(Nx.greater(x, 0), x, x * opts[:alpha])
   end
@@ -347,7 +348,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn linear(x), do: x
+  defblock linear(x), do: x
 
   @doc ~S"""
   Logsumexp activation.
@@ -372,24 +373,12 @@ defmodule Axon.Activations do
       >
 
   """
-  defn log_sumexp(x, opts \\ []) do
+  defblock LogSumExp, log_sumexp(x, opts \\ []) do
     opts = keyword!(opts, axis: -1)
     axes = wrap(opts[:axis])
 
-    # This is a scaling term designed to prevent over/under flow when x is very
-    # large. Consider cases where the intermediate value e^x with large positive
-    # x, e^x tends towards infinity or 0. This poisons the rest of the
-    # calculation which would otherwise be normalized with the division by sum(e^x).
-    # Thus we can scale by the max value in the tensor which guarantees all values
-    # are smaller than 0.
-    #
-    # Given the expression is essentially:
-    #
-    # e^(x - C) / sum(e^(x - C))
-    #
-    # We are essentially treating the max value as a constant term, C. Thus there
-    # is no need to differentiate through the max. See also: https://github.com/google/jax/pull/2260
-    # for a note on performance.
+    # Scaling term to prevent over/underflow; max treated as constant C.
+    # See also: https://github.com/google/jax/pull/2260
     max_val = Nx.reduce_max(x, axes: axes, keep_axes: true)
     max_val = stop_grad(Nx.select(Nx.is_infinity(max_val), 0, max_val))
 
@@ -398,13 +387,10 @@ defmodule Axon.Activations do
       |> Nx.subtract(max_val)
       |> Nx.exp()
 
-    res =
-      stable_exp
-      |> Nx.sum(axes: axes, keep_axes: true)
-      |> Nx.log()
-      |> Nx.add(max_val)
-
-    res
+    stable_exp
+    |> Nx.sum(axes: axes, keep_axes: true)
+    |> Nx.log()
+    |> Nx.add(max_val)
   end
 
   @doc ~S"""
@@ -430,7 +416,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn log_sigmoid(x), do: -softplus(-x)
+  defblock log_sigmoid(x), do: -softplus(-x)
 
   @doc """
   Log-softmax activation.
@@ -454,7 +440,7 @@ defmodule Axon.Activations do
         ]
       >
   """
-  defn log_softmax(x, opts \\ []) do
+  defblock LogSoftMax, log_softmax(x, opts \\ []) do
     opts = keyword!(opts, axis: -1)
 
     shifted = x - stop_grad(Nx.reduce_max(x, axes: [opts[:axis]], keep_axes: true))
@@ -489,7 +475,7 @@ defmodule Axon.Activations do
         ]
       >
   """
-  defn mish(x) do
+  defblock mish(x) do
     x * tanh(softplus(x))
   end
 
@@ -516,7 +502,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn relu(x) do
+  defblock ReLU, relu(x) do
     custom_grad(
       Nx.max(x, 0),
       [x],
@@ -551,7 +537,7 @@ defmodule Axon.Activations do
     * [MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications](https://arxiv.org/abs/1704.04861v1)
 
   """
-  defn relu6(x) do
+  defblock ReLU6, relu6(x) do
     x
     |> Nx.max(0)
     |> Nx.min(6)
@@ -586,9 +572,13 @@ defmodule Axon.Activations do
 
   """
   defn sigmoid(x) do
-    # Cache logits so they are available in certain calculations,
-    # e.g. binary_cross_entropy and categorical_cross_entropy
-    cache_logits(x, Nx.sigmoid(x))
+    # Logits metadata must wrap the block result (not sit inside it) so
+    # losses can pattern-match `%{logits: _}` on the returned tensor.
+    cache_logits(x, sigmoid_block(x))
+  end
+
+  defblockp Sigmoid, sigmoid_block(x) do
+    Nx.sigmoid(x)
   end
 
   @doc ~S"""
@@ -618,7 +608,7 @@ defmodule Axon.Activations do
     * [Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning](https://arxiv.org/abs/1702.03118v3)
 
   """
-  defn silu(x) do
+  defblock SiLU, silu(x) do
     x
     |> Nx.sigmoid()
     |> Nx.multiply(x)
@@ -655,7 +645,7 @@ defmodule Axon.Activations do
     * [Self-Normalizing Neural Networks](https://arxiv.org/abs/1706.02515v5)
 
   """
-  defn selu(x, opts \\ []) do
+  defblock SeLU, selu(x, opts \\ []) do
     opts =
       keyword!(opts,
         alpha: 1.6732632423543772848170429916717,
@@ -701,39 +691,28 @@ defmodule Axon.Activations do
 
   """
   defn softmax(x, opts \\ []) do
+    # Logits metadata must wrap the block result (not sit inside it) so
+    # losses can pattern-match `%{logits: _}` on the returned tensor.
+    cache_logits(x, softmax_block(x, opts))
+  end
+
+  defblockp SoftMax, softmax_block(x, opts \\ []) do
     opts = keyword!(opts, axis: -1)
     axes = wrap(opts[:axis])
 
-    # This is a scaling term designed to prevent over/under flow when x is very
-    # large. Consider cases where the intermediate value e^x with large positive
-    # x, e^x tends towards infinity or 0. This poisons the rest of the
-    # calculation which would otherwise be normalized with the division by sum(e^x).
-    # Thus we can scale by the max value in the tensor which guarantees all values
-    # are smaller than 0.
-    #
-    # Given the expression is essentially:
-    #
-    # e^(x - C) / sum(e^(x - C))
-    #
-    # We are essentially treating the max value as a constant term, C. Thus there
-    # is no need to differentiate through the max. See also: https://github.com/google/jax/pull/2260
-    # for a note on performance.
+    # Scaling term to prevent over/underflow; max treated as constant C.
+    # See also: https://github.com/google/jax/pull/2260
     max_val = stop_grad(Nx.reduce_max(x, axes: axes, keep_axes: true))
 
     stable_exp =
       x
       |> Nx.subtract(max_val)
       |> Nx.exp()
 
-    res =
-      stable_exp
-      |> Nx.sum(axes: axes, keep_axes: true)
-      |> reciprocal()
-      |> Nx.multiply(stable_exp)
-
-    # Cache logits so they are available in certain calculations,
-    # e.g. binary_cross_entropy and categorical_cross_entropy
-    cache_logits(x, res)
+    stable_exp
+    |> Nx.sum(axes: axes, keep_axes: true)
+    |> reciprocal()
+    |> Nx.multiply(stable_exp)
   end
 
   @doc ~S"""
@@ -759,7 +738,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn softplus(x) do
+  defblock SoftPlus, softplus(x) do
     stable = Nx.max(0.0, x)
 
     x
@@ -793,7 +772,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn softsign(x) do
+  defblock SoftSign, softsign(x) do
     x
     |> Nx.abs()
     |> Nx.add(1)
@@ -824,7 +803,7 @@ defmodule Axon.Activations do
       >
 
   """
-  defn tanh(x), do: Nx.tanh(x)
+  defblock tanh(x), do: Nx.tanh(x)
 
   ## Helpers