From a656f8ec2baf0b09cd92c68c79e131814b0e8439 Mon Sep 17 00:00:00 2001
From: Sean Moriarity <smoriarity.5@gmail.com>
Date: Tue, 26 May 2026 20:22:01 -0400
Subject: [PATCH] Add a residual layer combinator

---
 lib/axon.ex | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/lib/axon.ex b/lib/axon.ex
index e2e5eda3..1145c9dd 100644
--- a/lib/axon.ex
+++ b/lib/axon.ex
@@ -2509,6 +2509,43 @@ defmodule Axon do
     def unquote(op)([%Axon{} | _] = inputs), do: unquote(op)(inputs, [])
   end
 
+  @doc """
+  Adds a residual connection to the network.
+
+  Computes `fun.(x) + transform.(x)`. This is a convenience equivalent
+  for Axon.add(fun.(x), transform.(x)).
+
+  `fun` is the residual branch applied to `x`, typically a small
+  subnetwork. `transform` is applied to `x` along the skip path
+  before the addition and defaults to the identity, producing a
+  plain skip connection. Use `transform` to project the skip path
+  when `fun.(x)` and `x` have incompatible shapes (e.g. a 1x1
+  convolution, a dense projection, or a downsampling step).
+
+  ## Options
+
+    * `:name` - layer name.
+
+  ## Examples
+
+      iex> x = Axon.input("input", shape: {nil, 64})
+      iex> Axon.residual(x, &Axon.dense(&1, 64))
+
+      iex> x = Axon.input("input", shape: {nil, 32})
+      iex> Axon.residual(x, &Axon.dense(&1, 64), &Axon.dense(&1, 64))
+  """
+  @doc type: :combinator
+  def residual(%Axon{} = x, fun, transform \\ &Function.identity/1, opts \\ [])
+      when is_function(fun, 1) and is_function(transform, 1) and is_list(opts) do
+    opts = Keyword.validate!(opts, [:name, :meta])
+
+    layer(:add, [container({fun.(x), transform.(x)})],
+      name: opts[:name],
+      meta: opts[:meta],
+      op_name: :residual
+    )
+  end
+
   @doc """
   Adds a conditional layer which conditionally executes
   `true_graph` or `false_graph` based on the condition `cond_fn`