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