Add scale layer

lib/axon.ex

@@ -3425,6 +3425,47 @@ defmodule Axon do
     layer(:bias, [x, bias], name: opts[:name], meta: opts[:meta], op_name: :bias)
   end
 
+  @doc """
+  Adds a learnable per-channel scale layer to the network.
+
+  A scale layer multiplies the input elementwise by a trainable
+  parameter aligned to the input's `:channel_index` axis.
+
+  This is commonly used as the `gamma` in residual blocks of modern
+  Transformer architectures (CaiT, ConvNeXt, BEiT, EVA, etc.).
+
+  ## Options
+
+    * `:name` - layer name.
+
+    * `:scale_initializer` - initializer for the scale weights.
+      Defaults to `Axon.Initializers.full(1.0e-6)`.
+
+    * `:channel_index` - input feature axis along which the scale is
+      broadcast. Defaults to `-1`.
+  """
+  @doc type: :linear
+  def scale(%Axon{} = x, opts \\ []) do
+    opts =
+      Keyword.validate!(opts, [
+        :name,
+        :meta,
+        scale_initializer: Axon.Initializers.full(1.0e-6),
+        channel_index: -1
+      ])
+
+    channel_index = opts[:channel_index]
+
+    scale = param("scale", [{:axis, channel_index}], initializer: opts[:scale_initializer])
+
+    layer(:scale, [x, scale],
+      name: opts[:name],
+      meta: opts[:meta],
+      channel_index: channel_index,
+      op_name: :scale
+    )
+  end
+
   @doc """
   Adds a stack columns layer to the network.
 

lib/axon/layers.ex

@@ -2001,6 +2001,32 @@ defmodule Axon.Layers do
     input + bias
   end
 
+  @doc ~S"""
+  Functional implementation of a learnable scale layer.
+
+  Multiplies the input elementwise by `scale` broadcast along the
+  `:channel_index` axis:
+
+  $$y = x * \gamma$$
+
+  where $\gamma$ is the scale parameter with one element per
+  position along `:channel_index`.
+
+  ## Options
+
+    * `:channel_index` - input axis along which the scale is
+      applied. Defaults to `-1`.
+  """
+  defn scale(input, scale, opts \\ []) do
+    opts = keyword!(opts, channel_index: -1, mode: :train)
+
+    num_channels = Nx.axis_size(input, opts[:channel_index])
+    parameter_shape = norm_parameter_reshape(input, num_channels, opts[:channel_index])
+
+    scale = Nx.reshape(scale, parameter_shape)
+    Nx.multiply(input, scale)
+  end
+
   @doc """
   Resizes a batch of tensors to the given shape using one of a
   number of sampling methods.

test/axon/compiler_test.exs

@@ -347,6 +347,70 @@ defmodule CompilerTest do
     end
   end
 
+  describe "scale" do
+    test "initializes in default case" do
+      model = Axon.input("input_0", shape: {nil, 3}) |> Axon.scale(name: "scale")
+
+      input = random({1, 3})
+
+      assert {init_fn, _predict_fn} = Axon.build(model)
+
+      assert %ModelState{
+               data: %{"scale" => %{"scale" => scale}},
+               parameters: %{"scale" => ["scale"]}
+             } = init_fn.(input, ModelState.empty())
+
+      assert Nx.shape(scale) == {3}
+      assert Nx.type(scale) == {:f, 32}
+    end
+
+    test "applies small init multiplicatively" do
+      model =
+        Axon.input("input_0", shape: {nil, 3})
+        |> Axon.scale(name: "scale", scale_initializer: Axon.Initializers.full(1.0e-6))
+
+      input = Nx.tensor([[1.0, 2.0, 3.0]])
+
+      assert {init_fn, predict_fn} = Axon.build(model)
+      params = init_fn.(input, ModelState.empty())
+
+      out = predict_fn.(params, input)
+      assert_all_close(out, Nx.tensor([[1.0e-6, 2.0e-6, 3.0e-6]]), atol: 1.0e-12)
+    end
+
+    test "broadcasts correctly along a non-trailing channel_index" do
+      model =
+        Axon.input("input_0", shape: {nil, 2, 4})
+        |> Axon.scale(name: "scale", channel_index: 1)
+
+      input =
+        Nx.tensor([
+          [
+            [1.0, 1.0, 1.0, 1.0],
+            [1.0, 1.0, 1.0, 1.0]
+          ]
+        ])
+
+      assert {init_fn, predict_fn} = Axon.build(model)
+      params = init_fn.(input, ModelState.empty())
+
+      params =
+        Axon.ModelState.update(params, %{"scale" => %{"scale" => Nx.tensor([2.0, 3.0])}})
+
+      out = predict_fn.(params, input)
+
+      expected =
+        Nx.tensor([
+          [
+            [2.0, 2.0, 2.0, 2.0],
+            [3.0, 3.0, 3.0, 3.0]
+          ]
+        ])
+
+      assert_equal(out, expected)
+    end
+  end
+
   describe "dense" do
     test "initializes in default case" do
       model = Axon.input("input_0", shape: {nil, 1}) |> Axon.dense(1, name: "dense")