Add SwiGLU activation

lib/axon/activations.ex

@@ -780,6 +780,60 @@ defmodule Axon.Activations do
     |> Nx.multiply(x)
   end
 
+  @doc ~S"""
+  Swish-gated linear unit activation.
+
+  SwiGLU splits the input tensor along the given axis into two equal
+  halves $a$ and $b$, then returns $silu(a) \odot b$. The dimension of
+  the input along the given axis must be divisible by 2.
+
+  $$f(x) = silu(a) \odot b \quad \text{where} \quad x = [a, b]$$
+
+  ## Options
+
+    * `:axis` - axis along which to split the input into the activation
+      and gate halves. Defaults to `-1`.
+
+  ## Examples
+
+      iex> Axon.Activations.swiglu(Nx.tensor([[-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0]], names: [:batch, :data]))
+      #Nx.Tensor<
+        f32[batch: 1][data: 4]
+        [
+          [-0.14227763, -0.47681168, -0.8068243, 0.0]
+        ]
+      >
+
+  ### Error cases
+
+      iex> Axon.Activations.swiglu(Nx.tensor([1.0, 2.0, 3.0]))
+      ** (ArgumentError) axis -1 of input to swiglu must have a dimension divisible by 2, got dimension of size 3
+
+  ## References
+    * [GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202)
+  """
+  defblock SwiGLU, swiglu(x, opts \\ []) do
+    opts = keyword!(opts, axis: -1)
+    {a, b} = split_halves(x, opts[:axis])
+    silu(a) * b
+  end
+
+  deftransformp split_halves(x, axis) do
+    axis_idx = Nx.axis_index(x, axis)
+    size = elem(Nx.shape(x), axis_idx)
+
+    if rem(size, 2) != 0 do
+      raise ArgumentError,
+            "axis #{inspect(axis)} of input to swiglu must have a dimension" <>
+              " divisible by 2, got dimension of size #{size}"
+    end
+
+    half = div(size, 2)
+    a = Nx.slice_along_axis(x, 0, half, axis: axis_idx)
+    b = Nx.slice_along_axis(x, half, half, axis: axis_idx)
+    {a, b}
+  end
+
   @doc ~S"""
   Hyperbolic tangent activation.
 

lib/axon/layers.ex

@@ -2269,7 +2269,7 @@ defmodule Axon.Layers do
   # to use when invoking activation layers.
   @activation_layers [:exp, :gelu, :hard_tanh, :linear, :log_sigmoid] ++
                        [:mish, :relu, :relu6, :sigmoid, :silu, :softplus] ++
-                       [:softsign, :tanh]
+                       [:softsign, :tanh, :swiglu]
 
   for activation <- @activation_layers do
     @doc false