diff --git a/.tool-versions b/.tool-versions
new file mode 100644
index 00000000..1632f3f2
--- /dev/null
+++ b/.tool-versions
@@ -0,0 +1,2 @@
+erlang 27.3
+elixir 1.20.2-otp-27
diff --git a/lib/axon/activations.ex b/lib/axon/activations.ex
index dee3674d..954e6172 100644
--- a/lib/axon/activations.ex
+++ b/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,23 +691,17 @@ 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 =
@@ -725,15 +709,10 @@ defmodule Axon.Activations do
       |> 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
 
diff --git a/lib/axon/block.ex b/lib/axon/block.ex
new file mode 100644
index 00000000..6b6727e3
--- /dev/null
+++ b/lib/axon/block.ex
@@ -0,0 +1,201 @@
+defmodule Axon.Block do
+  @moduledoc """
+  Defines reusable `Nx.block/4` layers via the `defblock` / `defblockp` macros.
+
+  Both macros expand to:
+
+    * a struct module under the **caller module**, used as the `Nx.block/4` tag
+    * a private `defnp` with the layer body
+    * a `deftransform` (`defblock`) or `deftransformp` (`defblockp`) that wraps
+      that body in `Nx.block/4`
+
+  The body lives in `defnp` so `BinaryBackend.block/4` re-running the default
+  callback does not invoke `stop_grad`/`custom_grad` as raw Kernel calls on
+  concrete tensors. The callback instead calls the `defnp`, which JIT-compiles
+  normally and returns concrete results.
+
+  Use `defblockp` when the block is an implementation detail of a public
+  function (for example `softmax` wrapping `softmax_block`) so the module does
+  not advertise the block entry point.
+
+  By default the struct module is `CallerModule.<CamelizedFunName>`:
+
+      defblock selu(x, opts \\\\ []) do
+        ...
+      end
+
+  defined in `Axon.Activations` yields `%Axon.Activations.Selu{}`.
+
+  Pass an optional single-segment alias when you need non-default casing:
+
+      defblock SeLU, selu(x, opts \\\\ []) do
+        ...
+      end
+
+  yields `%Axon.Activations.SeLU{}`.
+
+  Trailing keyword `opts \\\\ []` (or any list default) arguments are stored on the
+  block struct as `:opts` and are **not** passed in the `Nx.block/4` args list.
+  That matches current Nx: block args must be tensors (or containers of tensors);
+  static options live on the struct. The block lambda restores `opts` from the
+  struct before calling the private `defnp`, so bodies can call `keyword!/2`
+  unchanged.
+
+  The struct is the dispatch tag for custom kernel implementations
+  (for example `defimpl EXLA.CustomCall, for: Axon.Activations.ReLU`).
+  It remains defined even when using `defblockp`.
+
+  The module that calls `defblock`/`defblockp` must `import Nx.Defn` so the
+  generated definitions are in scope.
+
+  ## Examples
+
+      defmodule MyLayers do
+        import Nx.Defn
+        import Axon.Block
+
+        defblock dense(x, w, b) do
+          x |> Nx.dot(w) |> Nx.add(b)
+        end
+      end
+
+  This defines `MyLayers.dense/3` and the struct `%MyLayers.Dense{}`.
+
+      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
+  """
+
+  @doc """
+  Defines a public block under the caller module, camelizing the function name.
+  """
+  defmacro defblock(call, do: body) do
+    build(__CALLER__, nil, call, body, :deftransform)
+  end
+
+  @doc """
+  Defines a public block under the caller module with an explicit module suffix.
+
+  `suffix` must be a single-segment alias (for example `SeLU`), not a nested
+  module path.
+  """
+  defmacro defblock(suffix, call, do: body) do
+    build(__CALLER__, suffix, call, body, :deftransform)
+  end
+
+  @doc """
+  Like `defblock/1`, but the wrapper is a private `deftransformp`.
+  """
+  defmacro defblockp(call, do: body) do
+    build(__CALLER__, nil, call, body, :deftransformp)
+  end
+
+  @doc """
+  Like `defblock/2`, but the wrapper is a private `deftransformp`.
+  """
+  defmacro defblockp(suffix, call, do: body) do
+    build(__CALLER__, suffix, call, body, :deftransformp)
+  end
+
+  defp build(env, suffix_ast, call, body, kind) do
+    {name, args} = parse_call(call, env)
+    {tensor_args, opts_args} = split_opts_args(args)
+    tensor_vars = Enum.map(tensor_args, &arg_var/1)
+    defn_args = Enum.map(args, &arg_var/1)
+    struct_module = Module.concat(env.module, suffix_name(suffix_ast, name, env))
+    defn_name = :"__block__#{name}__"
+
+    {struct_def, struct} =
+      case opts_args do
+        [] ->
+          {
+            quote(do: defstruct([])),
+            quote(do: %unquote(struct_module){})
+          }
+
+        [opts_arg] ->
+          opts_var = arg_var(opts_arg)
+
+          {
+            quote(do: defstruct(opts: [])),
+            quote(do: %unquote(struct_module){opts: unquote(opts_var)})
+          }
+
+        other ->
+          raise CompileError,
+            description:
+              "defblock/defblockp supports at most one trailing opts argument, got: " <>
+                Macro.to_string(other),
+            file: env.file,
+            line: env.line
+      end
+
+    quote do
+      defmodule unquote(struct_module) do
+        @moduledoc false
+        unquote(struct_def)
+      end
+
+      # Transform first so a preceding @doc attaches here, not to defnp.
+      unquote(kind)(unquote(name)(unquote_splicing(args))) do
+        Nx.block(
+          unquote(struct),
+          [unquote_splicing(tensor_vars)],
+          nil,
+          fn unquote(struct), unquote_splicing(tensor_vars) ->
+            unquote(defn_name)(unquote_splicing(defn_args))
+          end
+        )
+      end
+
+      defnp unquote(defn_name)(unquote_splicing(args)) do
+        unquote(body)
+      end
+    end
+  end
+
+  defp split_opts_args(args) do
+    Enum.split_while(args, fn arg -> not opts_arg?(arg) end)
+  end
+
+  defp opts_arg?({:\\, _meta, [_var, default]}), do: is_list(default)
+  defp opts_arg?(_), do: false
+
+  defp suffix_name(nil, name, _env) do
+    name |> Atom.to_string() |> Macro.camelize()
+  end
+
+  defp suffix_name({:__aliases__, _meta, [segment]}, _name, _env) when is_atom(segment) do
+    Atom.to_string(segment)
+  end
+
+  defp suffix_name(other, _name, env) do
+    raise CompileError,
+      description:
+        "defblock/defblockp optional name must be a single-segment alias like `SeLU`, got: " <>
+          Macro.to_string(other),
+      file: env.file,
+      line: env.line
+  end
+
+  defp parse_call({name, _meta, args}, _env) when is_atom(name) and is_list(args) do
+    {name, args}
+  end
+
+  defp parse_call({name, _meta, nil}, _env) when is_atom(name) do
+    {name, []}
+  end
+
+  defp parse_call(other, env) do
+    raise CompileError,
+      description:
+        "defblock/defblockp expects a function head like `name(args...)`, got: " <>
+          Macro.to_string(other),
+      file: env.file,
+      line: env.line
+  end
+
+  defp arg_var({:\\, _meta, [var, _default]}), do: var
+  defp arg_var(var), do: var
+end