Annotate activation functions with Nx blocks - #639
Conversation
Nest block structs under the caller module (defblock selu -> Activations.Selu, optional SeLU casing), keep opts on the struct so block args stay tensors, and move stop_grad/custom_grad/cache_logits outside block bodies so BinaryBackend evaluation does not leak Expr metadata. Update activation doctests for Nx 0.13 inspect formatting.
| defn relu(x) do | ||
| # custom_grad must sit outside Nx.block/4 — BinaryBackend.block re-runs the | ||
| # callback and Expr.metadata would otherwise leak into the result. | ||
| custom_grad( | ||
| Nx.max(x, 0), | ||
| relu_block(x), | ||
| [x], | ||
| fn g -> [Nx.select(Nx.greater(x, 0), g, Nx.broadcast(0, g))] end | ||
| ) | ||
| end | ||
|
|
||
| defblock Relu, relu_block(x) do | ||
| Nx.max(x, 0) | ||
| end |
There was a problem hiding this comment.
I'm not sure that custom_grad should sit outside of Nx.block. The grad of block is the grad of the default callback IIRC, so I don't think this should matter
There was a problem hiding this comment.
I'm not sure that custom_grad should sit outside of Nx.block. The grad of block is the grad of the default callback IIRC, so I don't think this should matter
You're right on the default callback, The thing I'm having trouble with is on forwarding. The callback gets re-run on real tensors and custom_grad/stop_grad turns results into a traced tensor
There was a problem hiding this comment.
custom_grad and stop_grad should work though. Maybe defblock is not using proper defn scoping for the callback body?
There was a problem hiding this comment.
maybe we should make sure that custom_gradand stop_grad are guarded on Nx?
Public API is a deftransform that builds Nx.block; the callback calls a private defnp so custom_grad/stop_grad JIT normally instead of wrapping concrete tensors. Pin nx/exla/torchx to elixir-nx/nx main for Nx.block.
| cache_logits(x, softmax_block(x, opts)) | ||
| end | ||
|
|
||
| defblock SoftMax, softmax_block(x, opts \\ []) do |
There was a problem hiding this comment.
Maybe this warrants defblockp where the inner deftransform is also defined with deftransformp so that this module just advertises softmax but not softmax_block
| end | ||
|
|
||
| # Public transform first so a preceding @doc attaches here, not to defnp. | ||
| deftransform unquote(name)(unquote_splicing(args)) do |
There was a problem hiding this comment.
| deftransform unquote(name)(unquote_splicing(args)) do | |
| unquote(kind)(unquote(name)(unquote_splicing(args))) do |
where kind is either :deftransform or :deftransformp coming from defblock and defblockp respectively, I think is the correct solution to my other comment in activations
Polvalente wanted the SoftMax/Sigmoid helpers hidden from the module API; defblockp generates deftransformp while leaving the block struct public.
The merge left deps/0 calling nx_opts/exla_opts/torchx_opts while the helpers were still the old github-pin *_dep functions, which broke CI at mix.exs compile. Hex already ships Nx.block in 0.13, so keep main's path-override helpers.
3 participants
Summary