Reference
The API reference of NNlib.
Activation Functions
Non-linearities that go between layers of your model. Note that, unless otherwise stated, activation functions operate on scalars. To apply them to an array you can call σ.(xs), relu.(xs) and so on.
For Float16 and BFloat16 inputs, the activations that involve arithmetic are computed in Float32 and the result is rounded back to the input type, as PyTorch does. This is faster and more accurate than half precision arithmetic.
NNlib.celu — Function
celu(x, α=1) = x ≥ 0 ? x : α * (exp(x/α) - 1)Activation function from "Continuously Differentiable Exponential Linear Units".
julia> lineplot(-2, 2, celu, height=7)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉│ celu(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠉⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⡧⠶⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⠤⠔⠒⠋⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⠤⠤⠤⠤⠔⠒⠒⠒⠊⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> celu(-10f0)
-0.9999546f0NNlib.elu — Function
elu(x, α=1) = x > 0 ? x : α * (exp(x) - 1)Exponential Linear Unit activation function. See "Fast and Accurate Deep Network Learning by Exponential Linear Units". You can also specify the coefficient explicitly, e.g. elu(x, 1).
julia> lineplot(-2, 2, elu, height=7)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉│ elu(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠉⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⡧⠶⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⠤⠔⠒⠋⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⠤⠤⠤⠤⠔⠒⠒⠒⠊⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> elu(-10f0)
-0.9999546f0
julia> elu(-10f0, 2)
-1.9999092f0NNlib.gelu — Function
gelu(x) = gelu_tanh(x)Activation function from "Gaussian Error Linear Units". See gelu_tanh.
NNlib.gelu_tanh — Function
gelu_tanh(x) = 0.5x * (1 + tanh(√(2/π) * (x + 0.044715x^3)))Activation function from "Gaussian Error Linear Units" using tanh approximation.
This implementation uses tanh which allows for better pattern matching and fusion in optimizing compilers compared to the sigmoid-based implementation. For a potentially faster implementation that uses sigmoid_fast, see gelu_sigmoid.
julia> lineplot(-2, 2, gelu_tanh, height=7)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊│ gelu_tanh(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠊⠁⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⣀⡠⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⣤⣤⣤⣤⣤⣤⣤⣤⡤⠤⠤⠤⠤⠤⠤⠤⣤⣤⣤⡤⡧⠶⠶⠭⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠈⠉⠉⠉⠉⠉⠉⠉⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> lineplot(-5, 0, gelu_tanh, height=7);
julia> lineplot!(ans, swish)
┌────────────────────────────────────────┐
0 │⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠒⠒⠤⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸│ gelu_tanh(x)
│⠑⠒⠢⠤⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇│ swish(x)
│⠀⠀⠀⠀⠀⠈⠉⠒⠤⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⢆⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣸⠁│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠒⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⢄⠀⠀⠀⠀⠀⠀⠀⠀⢠⡇⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠓⣄⠀⠀⠀⠀⠀⢠⡞⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠓⢄⣀⣀⡤⢣⠃⠀⠀│
-0.2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢠⠇⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀0⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.gelu_sigmoid — Function
gelu_sigmoid(x) = x * σ(√(8/π) * (x + 0.044715x^3))Alternative implementation of the GELU activation function using sigmoid instead of tanh. This is mathematically equivalent to gelu_tanh but may be faster in some cases.
The sigmoid-based implementation may prevent pattern matching and fusion in some optimizing compilers. Use gelu_tanh if you need better compiler optimization support.
Not to be confused with quick_gelu, the approximation x * σ(1.702x) used by CLIP.
NNlib.gelu_erf — Function
gelu_erf(x) = xΦ(x) = 0.5x * (1 + erf(x/√2))Activation function from "Gaussian Error Linear Units" without approximation. The SpecialFunctions.jl package needs to be loaded to use this function.
NNlib.quick_gelu — Function
quick_gelu(x) = x * σ(1.702x)Sigmoid approximation of GELU from "Gaussian Error Linear Units", called QuickGELU. It is the activation of OpenAI's CLIP, and of the CLIP models trained with it: MetaCLIP, DFN, and those named *quickgelu* in open_clip and timm.
This is a different function from gelu_sigmoid, which is the tanh approximation written with a sigmoid. The two differ by up to 0.021, near x = -2.3.
julia> lineplot(-2, 2, quick_gelu, height=7)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊│ quick_gelu(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠊⠁⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⣀⡠⠔⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⣤⣤⣤⣤⣤⣤⣤⣤⡤⠤⠤⠤⠤⠤⠤⢤⣤⣤⣤⡤⡧⠶⠶⠭⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠈⠉⠉⠉⠉⠉⠉⠁⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> quick_gelu(1f0), gelu_sigmoid(1f0)
(0.84579575f0, 0.841192f0)NNlib.hardsigmoid — Function
hardσ(x) = max(0, min(1, (x + 3) / 6))Piecewise linear approximation of sigmoid.
julia> lineplot(-5, 5, hardsigmoid, height=7)
┌────────────────────────────────────────┐
1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡠⠖⠋⠉⠉⠉⠉⠉⠉⠉⠉│ hardσ(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⣀⡤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⡠⠔⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⡗⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠊⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠋⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⠤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> lineplot(-5, 5, sigmoid, height=7)
┌────────────────────────────────────────┐
1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡠⠤⠖⠒⠒⠋⠉⠉⠉⠉⠉⠉│ σ(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣀⠔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⡏⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡔⠋⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠊⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⣀⣀⣀⣀⣀⣀⣀⠤⠤⠤⠒⠊⠉⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.sigmoid_fast — Function
sigmoid_fast(x)This used to be a faster, and very slightly less accurate, version of sigmoid. For Float32 and Float64 the two are now computed in the same way, and differ only in that sigmoid_fast(x) is exactly zero for x < -80.
See also tanh_fast.
julia> sigmoid(0.2f0)
0.54983395f0
julia> sigmoid_fast(0.2f0)
0.54983395f0
julia> hardσ(0.2f0)
0.53333336f0NNlib.hardtanh — Function
hardtanh(x) = max(-1, min(1, x))Segment-wise linear approximation of tanh, much cheaper to compute. See "Large Scale Machine Learning".
See also tanh_fast.
julia> lineplot(-2, 2, hardtanh, height=7)
┌────────────────────────────────────────┐
1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⠔⠋⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉│ hardtanh(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⣀⡤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⢀⡤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⡷⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠖⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠖⠋⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⠔⠋⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x
julia> lineplot(-2, 2, tanh, height=7)
┌────────────────────────────────────────┐
1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⣀⡠⠤⠤⠒⠒⠒⠊⠉⠉⠉│ tanh(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⢀⡠⠔⠊⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⢀⡤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⡷⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡠⠔⠊⠁⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⣀⣀⣀⡠⠤⠤⠤⠖⠒⠊⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.tanh_fast — Function
tanh_fast(x)This is a faster but slighly less accurate version of tanh.
Where Julia's tanh function has an error under 2 eps, this may be wrong by 5 eps, a reduction by less than one decimal digit.
For x::Float32 this is usually about 10 times faster, with a smaller speedup for x::Float64. Float16 and BFloat16 are computed in Float32. For any other number types, it just calls tanh.
See also sigmoid_fast.
julia> tanh(0.5f0)
0.46211717f0
julia> tanh_fast(0.5f0)
0.46211714f0
julia> hard_tanh(0.5f0)
0.5f0NNlib.leakyrelu — Function
leakyrelu(x, a=0.01) = max(a*x, x)Leaky Rectified Linear Unit activation function. You can also specify the coefficient explicitly, e.g. leakyrelu(x, 0.01).
julia> lineplot(-2, 2, x -> leakyrelu(x, 0.5), height=7)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉│ #42(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠉⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⣤⣤⡤⡧⠶⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⠤⠤⠒⠒⠋⠉⠁⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⣀⣀⠤⠤⠒⠒⠊⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> leakyrelu(-10f0, 0.2)
-2.0f0
julia> leakyrelu(-10f0, 0.02)
-0.5f0NNlib.lisht — Function
lisht(x) = x * tanh(x)Activation function from "LiSHT: Non-Parametric Linearly Scaled Hyperbolic Tangent ..."
julia> lineplot(-2, 2, lisht, height=7)
┌────────────────────────────────────────┐
2 │⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔│ lisht(x)
│⠀⠈⠑⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠊⠁⠀│
│⠀⠀⠀⠀⠈⠣⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠁⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠑⢆⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠊⠁⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠢⡄⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⠔⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⢄⡀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⢀⡠⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⠦⣄⣀⣀⣇⣀⣀⠤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> lineplot!(ans, logcosh)
┌────────────────────────────────────────┐
2 │⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔│ lisht(x)
│⠀⠈⠑⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠊⠁⠀│ logcosh(x)
│⠢⣄⠀⠀⠈⠣⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠁⠀⠀⣀⠔│
f(x) │⠀⠈⠑⠢⣀⠀⠀⠑⢆⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠊⠁⠀⣀⠔⠊⠁⠀│
│⠀⠀⠀⠀⠀⠉⠢⢄⡀⠉⠢⡄⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⠔⠋⠀⡠⠔⠋⠁⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⠦⣌⡓⢄⡀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⢀⡠⠖⣁⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⠪⠷⣦⣄⣀⣀⣇⣀⣀⣤⠶⠕⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.logcosh — Function
logcosh(x)Return log(cosh(x)) which is computed in a numerically stable way.
julia> lineplot(-5, 5, logcosh, height=7)
┌────────────────────────────────────────┐
5 │⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ logcosh(x)
│⠉⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠋│
│⠀⠀⠀⠑⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠊⠁⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠑⠦⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠊⠁⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⠦⡀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⠦⡀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⢀⡤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠑⠢⢄⣀⣀⣇⣀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.logsigmoid — Function
logσ(x)Return log(σ(x)) which is computed in a numerically stable way.
julia> lineplot(-5, 5, logsigmoid, height=7)
┌────────────────────────────────────────┐
0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡧⠤⠔⠒⠒⠒⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉│ logσ(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠉⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⢀⡤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⣀⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⡤⠖⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-6 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.mish — Function
mish(x) = x * tanh(softplus(x))Activation function from "Mish: A Self Regularized Non-Monotonic Neural Activation Function".
julia> lineplot(-5, 5, mish, height=7)
┌────────────────────────────────────────┐
5 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠖⠋│ mish(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠒⠁⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠔⠋⠁⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⢀⡠⠖⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡤⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣧⣔⣊⣁⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀│
-1 │⠀⠀⠀⠀⠀⠀⠀⠀⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.relu — Function
relu(x) = max(0, x)Rectified Linear Unit activation function.
julia> lineplot(-2, 2, relu, height=7)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠋│ relu(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠊⠁⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠊⠁⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⡠⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⡠⠖⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣇⠔⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.relu6 — Function
relu6(x) = min(max(0, x), 6)Rectified Linear Unit activation function capped at 6. See "Convolutional Deep Belief Networks" from CIFAR-10.
julia> lineplot(-10, 10, relu6, height=7)
┌────────────────────────────────────────┐
6 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠎⠉⠉⠉⠉⠉⠉⠉⠉│ relu6(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⡤⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⡠⠎⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⡔⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⡧⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-10⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀10⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.rrelu — Function
rrelu(x, lo=1/8, hi=1/3) = max(a*x, x)
# where `a` is randomly sampled from uniform distribution `U(lo, hi)`Randomized Leaky Rectified Linear Unit activation function. See "Empirical Evaluation of Rectified Activations" You can also specify the bound explicitly, e.g. rrelu(x, 0.0, 1.0).
julia> lineplot(-20, 10, rrelu, height=7)
┌────────────────────────────────────────┐
10 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠋│ rrelu(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⢀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⠤⣤⣤⢤⣤⣤⠤⠤⠤⢼⠮⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⣰⢀⣆⡄⣄⡄⡠⡰⠦⠷⡜⢢⠷⠳⠢⠊⠉⠉⠀⠀⠁⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠃⠉⠙⠘⠃⠈⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-10 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-20⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀10⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> extrema(rrelu.(fill(-10f0, 1000)))
(-3.3316886f0, -1.2548422f0)NNlib.selu — Function
selu(x) = λ * (x ≥ 0 ? x : α * (exp(x) - 1))
λ ≈ 1.05070...
α ≈ 1.67326...Scaled exponential linear units. See "Self-Normalizing Neural Networks".
julia> lineplot(-3, 2, selu, height=7)
┌────────────────────────────────────────┐
3 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ selu(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠤⠒│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⢀⣀⠤⠖⠊⠉⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣀⡠⠤⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⣉⠭⠛⡏⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⡤⠤⠒⠊⠉⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-2 │⠤⠤⠖⠒⠒⠒⠒⠒⠒⠒⠉⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> selu(-10f0)
-1.7580194f0NNlib.sigmoid — Function
σ(x) = 1 / (1 + exp(-x))Classic sigmoid activation function. Unicode σ can be entered as \sigma then tab, in many editors. The ascii name sigmoid is also exported.
See also sigmoid_fast.
julia> using UnicodePlots
julia> lineplot(-5, 5, sigmoid, height=7)
┌────────────────────────────────────────┐
1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡠⠤⠖⠒⠒⠋⠉⠉⠉⠉⠉⠉│ σ(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣀⠔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⡏⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡔⠋⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠊⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⣀⣀⣀⣀⣀⣀⣀⠤⠤⠤⠒⠊⠉⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> sigmoid === σ
trueNNlib.softplus — Function
softplus(x) = log(exp(x) + 1)See "Deep Sparse Rectifier Neural Networks", JMLR 2011.
julia> lineplot(-3, 3, softplus, height=7)
┌────────────────────────────────────────┐
4 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ softplus(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠁⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠔⠊⠁⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⣀⡠⠤⠒⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⡧⠤⠒⠊⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⣀⣀⣀⣀⣀⣀⣀⡠⠤⠤⠤⠤⠔⠒⠒⠚⠉⠉⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀3⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> lineplot!(ans, relu)
┌────────────────────────────────────────┐
4 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ softplus(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣠│ relu(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣠⡴⠞⠋⠁│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣤⡴⠞⠋⠁⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⣀⡠⢤⡲⠝⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⡧⠤⠒⠊⣉⠥⠚⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⣀⣀⣀⣀⣀⣀⣀⣠⣤⣤⣤⣤⣔⣒⣒⣚⣉⣉⣁⣀⣇⠴⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀3⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> softplus(16f0)
16.0f0NNlib.softshrink — Function
softshrink(x, λ=0.5) =
(x ≥ λ ? x - λ : (-λ ≥ x ? x + λ : 0))See "Softshrink Activation Function".
julia> lineplot(-2, 2, softshrink, height=7)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀│ softshrink(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠉⠁│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⣀⡤⠤⠒⠋⠁⠀⠀⠀⠀⠀⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⣤⡤⠤⠤⠤⠤⠤⠤⡧⠤⠤⠤⠤⠶⠮⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⢀⣀⠤⠖⠒⠉⠁⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⣀⠤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-2 │⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> lineplot!(ans, tanhshrink)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀│ softshrink(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⣉⡡│ tanhshrink(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⣀⡤⠤⣒⣋⠥⠤⠒⠊⠉⠁⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⣤⣤⣤⣤⡤⠤⠤⠤⠤⠤⠤⡷⠶⠶⠶⠶⠶⠾⠿⠯⠭⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⢀⣀⡠⠤⠖⢒⣋⠭⠗⠒⠉⠁⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠊⣉⠤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-2 │⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀
julia> softshrink.((-10f0, 10f0))
(-9.5f0, 9.5f0)NNlib.softsign — Function
softsign(x) = x / (1 + |x|)See "Quadratic Polynomials Learn Better Image Features" (2009).
julia> lineplot(-5, 5, softsign, height=7)
┌────────────────────────────────────────┐
1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⣀⣀⣀⠤⠤⠤⠤⠤│ softsign(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⣀⡤⠖⠒⠋⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⡔⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡯⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠁⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⣀⠤⠤⠒⠋⠁⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⠒⠒⠒⠒⠒⠊⠉⠉⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> lineplot!(ans, tanh)
┌────────────────────────────────────────┐
1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡤⠖⠊⠉⠉⠉⣉⣉⣉⣉⣉⠭⠭⠭⠭⠭│ softsign(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⡔⣃⡤⠖⠒⠋⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│ tanh(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣧⡞⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡯⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡴⠃⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⣀⠤⠤⠒⢋⠕⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⣒⣒⣒⣒⣒⣊⣉⣉⣉⣉⣁⣀⣀⡠⠤⠒⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> softsign(1f0)
0.5f0
julia> softsign(100f0)
0.990099f0NNlib.swish — Function
swish(x) = x * σ(x)Self-gated activation function. See "Swish: a Self-Gated Activation Function".
julia> lineplot(-2, 2, swish, height=7)
┌────────────────────────────────────────┐
2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤│ swish(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠋⠁⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⢀⣀⡤⠔⠊⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⣤⣤⡤⡧⠴⠶⠯⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠉⠑⠒⠒⠒⠒⠒⠒⠒⠒⠒⠒⠉⠉⠉⠉⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ NNlib.hardswish — Function
hardswish(x) = x * hardσ(x)Hard-Swish activation function. See "Searching for MobileNetV3".
julia> lineplot(-2, 5, hardswish, height = 7)
┌────────────────────────────────────────┐
5 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠔⠒⠉│ hardswish(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡠⠔⠒⠉⠁⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠖⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⢀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣇⣤⣤⣖⣚⣉⣁⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀│
-1 │⠉⠒⠒⠒⠒⠉⠉⠉⠉⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> lineplot(-4, 0, hardswish, height = 7);
julia> lineplot!(ans, swish)
┌────────────────────────────────────────┐
0 │⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⢣⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡜│ hardswish(x)
│⠒⠒⠢⠤⢄⣀⡀⠀⠀⠀⠀⠱⡄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⠎⠀│ swish(x)
│⠀⠀⠀⠀⠀⠀⠈⠉⠑⠒⠦⢄⣘⢄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡴⠃⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠑⡖⠦⢄⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⢔⠏⠁⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠣⣄⠀⠉⠑⠒⠦⠤⢄⣀⣀⣀⣀⡠⠤⠖⣊⠕⠁⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠓⠤⡀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠁⠀⠀⠀⠀⠀⠀⠀│
-0.4 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠉⠒⠢⠤⠤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-4⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀0⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> hardswish.(-5:5)'
1×11 adjoint(::Vector{Float64}) with eltype Float64:
-0.0 -0.0 -0.0 -0.333333 -0.333333 0.0 0.666667 1.66667 3.0 4.0 5.0NNlib.tanhshrink — Function
tanhshrink(x) = x - tanh(x)See "Tanhshrink Activation Function".
julia> lineplot(-3, 3, tanhshrink, height=7)
┌────────────────────────────────────────┐
3 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ tanhshrink(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡠⠤⠖⠊│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⣀⡠⠤⠒⠊⠉⠁⠀⠀⠀⠀│
f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⣤⡤⠤⠤⠤⠤⠤⠤⡷⠶⠶⠶⠶⠶⠮⠭⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│
│⠀⠀⠀⠀⠀⣀⡠⠴⠒⠊⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⡠⠴⠒⠊⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
-3 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀3⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
julia> tanhshrink.((-10f0, 10f0))
(-9.0f0, 9.0f0)NNlib.trelu — Function
trelu(x, theta=1) = x > theta ? x : 0Threshold gated rectified linear activation function. See "Zero-bias autoencoders and the benefits of co-adapting features"
julia> lineplot(-2, 4, trelu, height=7)
┌────────────────────────────────────────┐
4 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠋│ trelu(x)
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀│
f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠴⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⣠⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⡏⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣇⣀⣀⣀⣀⣀⣀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│
└────────────────────────────────────────┘
⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀4⠀
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ Attention
NNlib.scaled_dot_product_attention — Function
scaled_dot_product_attention(query, key, value, [bias]; [fdrop, mask, scale, is_causal, window, q_offset, q_lengths, kv_lengths, q_segments, kv_segments, maxlen])
sdpa(query, key, value, [bias]; kws...)Multihead dot product attention used in transformer architectures, with an explicit head axis.
The input arrays must have the shape (head_dim, nheads, seq_len, batch_size...), that is a feature (head) dimension, a heads dimension, a sequence-length dimension, then an arbitrary number of batch dimensions or none. The number of heads is taken from the second dimension of each input, so no nheads keyword is needed.
This is the reverse of the (batch_size, seq_len, nheads, head_dim) layout of JAX and Flax. It is not the reverse of PyTorch's (batch_size, nheads, seq_len, head_dim): an array ported from PyTorch by reversing its dimensions has size (head_dim, seq_len, nheads, batch_size) and still needs permutedims(x, (1, 3, 2, 4)). A missing permutation cannot be detected, since the array is a valid input whose heads are taken as the sequence.
Grouped-query attention (GQA) is supported: if key and value have fewer heads than query, they are shared across query heads. The number of query heads must then be divisible by the number of key/value heads.
Returns the attention output array of size (v_head_dim, nheads, q_len, batch_size...).
See also scaled_dot_product_attention_with_weights if you need the attention weights too.
Arguments
query: Query array of size(qk_head_dim, nheads, q_len, batch_size...).key: Key array of size(qk_head_dim, nkvheads, kv_len, batch_size...).value: Value array of size(v_head_dim, nkvheads, kv_len, batch_size...).bias: Eithernothingor an array broadcastable to size(kv_len, q_len, nheads, batch_size...). It will be added to the attention logits before applying the softmax. Defaultnothing.fdrop: A dropout function or layer to be applied on the attention weights right after the softmax. Defaultidentity(no dropout).mask: Eithernothingor a boolean array broadcastable to size(kv_len, q_len, nheads, batch_size...). The mask is applied to the attention logits just before the softmax. Seemake_causal_maskfor creating causal masks. Defaultnothing.scale: The multiplier used to scale thequeryᵀ * keydot products before the softmax. Defaultnothing, meaning1 / √(qk_head_dim).is_causal: Iftrue, a causal mask is applied so that each query position can only attend to key positions up to and including its own. Equivalent tomask = make_causal_mask(query; dims=3)for equal query and key lengths, but a fused implementation does not have to materialize the mask. If:bottom_right, the mask is aligned with the last valid key instead of the first: the queries are taken as the last positions of the keys, as when decoding with a KV cache, andq_offsetiskv_lengths .- q_lengths, these being the sizes of the arrays where not given. This is the causal mask of FlashAttention and PyTorch'scausal_lower_right. Defaultfalse.window: Eithernothingor a positive integer, the size of a sliding window: the queryiattends only to the keysjwith|q_offset + i - j| < window. Withis_causal, these are the lastwindowkeys up to its own position,q_offset + i - window < j ≤ q_offset + i, as in the sliding-window layers of Mistral, Gemma or gpt-oss. This is thesliding_windowof transformers, and thewindow_sizeof FlashAttention is(window - 1, 0)withis_causal,(window - 1, window - 1)without. A fused implementation only goes through the keys of the windows. Defaultnothing.q_offset: Withis_causal=true, the number of keys that precede the first query: the queryiattends to the keysj ≤ q_offset + i. An integer, or an integer array with one entry for each batch element. It counts the keys given, which with a sliding-window cache is not the absolute position. It has no effect withoutis_causalorwindow, and is set byis_causal=:bottom_right. Default0.q_lengths,kv_lengths: Eithernothingor the number of valid positions of each query and of each key/value sequence in the batch, the following ones being padding: an integer array with one entry for each batch element, or an integer for all of them. The keys pastkv_lengthsare not attended to, and the output is zero at the queries pastq_lengths. Each can be given without the other. Defaultnothing.q_segments,kv_segments: Eithernothingor the lengths of the sequences packed one after the other along the sequence dimension ofqueryand ofkeyandvalue, as an integer vector. See below.kv_segmentsdefaults toq_segments.maxlen: Withq_segments, eithernothingor an integer at least as large as every segment. It saves reading the lengths of the segments on the host where they are on a GPU, which waits for the device, and stands in for them where they cannot be read, as the traced ones of Reactant. Defaultnothing.
mask, is_causal, window and kv_lengths can be combined: a key is then attended to only where all of them allow it. A query that attends to no key, as a padding query past the window of every valid key, has a NaN output, unless q_lengths makes it zero.
Packed sequences
With q_segments, the inputs hold sequences of different lengths one after the other, without padding: they have size (head_dim, nheads, total_len), without batch dimensions, and the output has the size of query. The s-th segment of query attends to the s-th segment of key and value only, as a batch element does in a padded batch: is_causal and window apply within each segment, and q_offset, q_lengths and kv_lengths take one entry for each segment. The lengths of the segments have to add up to the sequence lengths of the arrays. A bias or a mask is not supported.
This is how documents packed in a sequence are kept from attending to each other, and how a batch of sequences of different lengths avoids padding: layers other than the attention act on the (features, total_len) array of all the positions.
Fused implementation
Without an fdrop, the output is computed tile by tile, with the softmax taken online over the tiles of keys: the attention weights are never materialized, so that memory grows linearly rather than quadratically with the sequence length. This is what Arrays and GPU arrays of Float16, Float32 and Float64 take, and on the GPU those of BFloat16, with any of the arguments above. The bias and the mask, if any, have to be on the device of the inputs. On the GPU, a call with few queries per sequence, as in decoding with a key-value cache, splits the keys of each sequence over several blocks, and reads them once for all the query heads sharing them.
Any other call takes the materialized implementation, which forms the weights. NNlib.is_fused_attention tells which of the two a call takes. With a window, it forms them over the band of the window only, for blocks of queries with the keys of their windows, so that memory grows linearly with the sequence length there too. fdrop is then given these blocks of weights rather than the whole array.
The fused implementation has a gradient with respect to query, key, value and scale, which recomputes the weights tile by tile and does not materialize them either. It is what automatic differentiation based on ChainRules (as that of Zygote) and the reverse mode of Enzyme use. Under them a call with a bias takes the materialized implementation, which gives the gradient with respect to the bias too, and so does every call under the other automatic differentiation packages. A query that attends to no key, whose output is NaN, takes no part in the fused gradient.
With CUDA.jl and cuDNN.jl loaded, the Float16 and BFloat16 CuArrays that cuDNN's own fused kernel has an engine for are handed to it: it runs on tensor cores, and has a gradient. It supports scale, is_causal, q_lengths, kv_lengths and grouped-query attention. cuDNN.jl 6.5 also supports bottom-right causal masks, sliding windows and packed sequences. It requires:
- no
bias,fdropormask(that isnothing,identityandnothing); - with an active causal mask, top-left or bottom-right alignment; without one, a window must be top-left aligned;
- a head dimension that is a multiple of 8 and at most 256;
- a
valuewith the same head dimension asquery.
Earlier cuDNN.jl versions require no window or packed sequences, and a top-left causal mask when it excludes any keys. Unsupported calls use NNlib's own implementation.
Examples
# (head_dim, nheads, seq_len, batch)
q = rand(Float32, 16, 8, 20, 4)
k = v = rand(Float32, 16, 8, 30, 4)
y = scaled_dot_product_attention(q, k, v)
# size(y) == (16, 8, 20, 4)
# grouped-query attention: 8 query heads sharing 2 key/value heads
q = rand(Float32, 16, 8, 20, 4)
k = v = rand(Float32, 16, 2, 30, 4)
y = scaled_dot_product_attention(q, k, v)
# a padded batch, whose first key/value sequence has 12 valid positions out of 30
y = scaled_dot_product_attention(q, k, v; kv_lengths=[12, 30, 30, 30])
# decoding with a KV cache of 64 positions: 3 new queries for each batch element,
# appended to the 10, 20, 30 and 40 keys already cached
q = rand(Float32, 16, 8, 3, 4)
k = v = rand(Float32, 16, 8, 64, 4)
y = scaled_dot_product_attention(q, k, v; is_causal=:bottom_right, kv_lengths=[13, 23, 33, 43])
# the same, with the offsets of the queries
y = scaled_dot_product_attention(q, k, v; is_causal=true, q_offset=[10, 20, 30, 40],
kv_lengths=[13, 23, 33, 43])
# sliding-window attention: each query attends to its last 8 keys, its own included
x = rand(Float32, 16, 8, 100, 4)
y = scaled_dot_product_attention(x, x, x; is_causal=true, window=8)
# three documents of 5, 12 and 3 positions packed in one sequence, each attending
# causally to itself only
x = rand(Float32, 16, 8, 20)
y = scaled_dot_product_attention(x, x, x; is_causal=true, q_segments=[5, 12, 3])
# size(y) == (16, 8, 20)NNlib.scaled_dot_product_attention_with_weights — Function
scaled_dot_product_attention_with_weights(query, key, value, [bias]; kws...)
sdpa_with_weights(query, key, value, [bias]; kws...)Return the output of scaled_dot_product_attention together with the attention weights it is computed from, in a single pass.
The arguments are those of scaled_dot_product_attention. The weights have size (kv_len, q_len, nheads, batch_size...), with α[j, i, h, b] the weight of the key j for the query i in the head h: the softmax over the keys of the attention logits, after fdrop. They are zero at the queries past q_lengths. With q_segments they have size (kv_len, q_len, nheads), and are zero but for the keys of the segment of each query.
Returning the weights means materializing them, so this function never takes a fused implementation: prefer scaled_dot_product_attention when they are not needed.
Examples
q = rand(Float32, 16, 8, 20, 4)
k = v = rand(Float32, 16, 8, 30, 4)
y, α = scaled_dot_product_attention_with_weights(q, k, v)
# size(y) == (16, 8, 20, 4)
# size(α) == (30, 20, 8, 4)NNlib.make_causal_mask — Function
make_causal_mask(x; dims)Return a boolean square matrix m of the same type as x and of side size(x, dims). Its elements are set such that m[i, j] == i ≤ j.
Can be used as the mask of scaled_dot_product_attention, whose inputs have the sequence length along dims=3. For a plain causal mask, prefer the is_causal=true keyword there, which lets a fused implementation avoid materializing the mask.
Calling make_causal_mask(x) without dims is deprecated. It keeps returning the mask for dims=2, the sequence dimension of the deprecated dot_product_attention.
NNlib.is_fused_attention — Function
NNlib.is_fused_attention(query, key, value, [bias]; kws...)Return true if scaled_dot_product_attention called with the same arguments is computed by a fused implementation, which does not materialize the attention weights, and false if it takes the materialized one.
This is about a call outside of automatic differentiation: under it, the materialized implementation is differentiated where the fused one has no gradient, as with a bias.
Rotary position embedding
NNlib.rotary_embedding_table — Function
NNlib.rotary_embedding_table(head_dim, maxlen, [T=Float32]; base=10_000)Return the tables (cos, sin) of the rotary position embedding (RoPE) of Su et al. for the positions 1:maxlen, to be passed to NNlib.apply_rotary_embedding.
Each is a Matrix{T} of size (head_dim ÷ 2, maxlen): the cosine and the sine of the angle θ[i] * (p - 1) by which the i-th pair of features is rotated at the position p, with the frequencies θ[i] = base^(-2(i - 1) / head_dim). The first position is not rotated, as the position zero of PyTorch and JAX.
The tables are computed in Float64 on the CPU: move them to the device of the inputs. Other frequencies, as those rescaled by Llama 3 or YaRN, give their own tables cos.(θ .* (0:maxlen - 1)') and sin.(θ .* (0:maxlen - 1)').
Examples
julia> cos_table, sin_table = NNlib.rotary_embedding_table(64, 4096);
julia> size(cos_table), size(sin_table)
((32, 4096), (32, 4096))NNlib.apply_rotary_embedding — Function
NNlib.apply_rotary_embedding(x, cos, sin; offset=0, positions=nothing, interleaved=false)Apply the rotary position embedding (RoPE) of Su et al. to x, a query or a key of scaled_dot_product_attention of size (head_dim, nheads, seq_len, batch_size...): the features of each head are rotated pair by pair, by angles proportional to the position of the token. The dot product of a query and a key so rotated depends on their positions through their difference only. The values are not rotated.
The tables cos and sin have size (head_dim ÷ 2, maxlen), their column p holding the cosines and the sines of the angles of the position p, as those of NNlib.rotary_embedding_table. At the position p, the k-th pair of features (a, b) becomes (a * cos[k, p] - b * sin[k, p], b * cos[k, p] + a * sin[k, p]). The pairs are (x[k], x[k + head_dim ÷ 2]), from the two halves of the features (as in Llama, Hugging Face Transformers and Flax), or with interleaved=true the consecutive features (x[2k - 1], x[2k]) (as in RoFormer, GPT-J and torchtune).
Returns an array of the size of x.
Keywords
offset: The tokeniis at the positionoffset + i. An integer, or an integer array with one entry for each batch element. With a KV cache, the number of tokens already in the cache, as theq_offsetofscaled_dot_product_attention. Default0.positions: Eithernothingor the positions of the tokens, in place ofoffset, as an integer array of size(seq_len, batch_size...), or(seq_len,)for all the batch elements. They are the columns of the tables. Defaultnothing.interleaved: Iftrue, the pairs are the consecutive features, else the two halves. Defaultfalse.
The positions have to be within the columns of the tables. To rotate part of the features only (as GPT-NeoX and Phi), rotate the slice x[1:d, :, :, :] with tables of d ÷ 2 rows and concatenate the rest.
Examples
The causal self-attention of a decoder, in which the queries and the keys are rotated but not the values, here with 8 heads of dimension 64 sharing 2 key/value heads:
using NNlib: apply_rotary_embedding, rotary_embedding_table
cos_table, sin_table = rotary_embedding_table(64, 4096)
q = rand(Float32, 64, 8, 100, 4) # (head_dim, nheads, seq_len, batch_size)
k = rand(Float32, 64, 2, 100, 4)
v = rand(Float32, 64, 2, 100, 4)
q = apply_rotary_embedding(q, cos_table, sin_table)
k = apply_rotary_embedding(k, cos_table, sin_table)
y = scaled_dot_product_attention(q, k, v; is_causal=true)Decoding with a KV cache, which holds the keys already rotated: the new tokens are rotated with the number of cached tokens as offset, the q_offset of the attention.
kcache = apply_rotary_embedding(rand(Float32, 64, 2, 10, 4), cos_table, sin_table)
vcache = rand(Float32, 64, 2, 10, 4)
# 3 new tokens at the positions 11 to 13
q = apply_rotary_embedding(rand(Float32, 64, 8, 3, 4), cos_table, sin_table; offset=10)
knew = apply_rotary_embedding(rand(Float32, 64, 2, 3, 4), cos_table, sin_table; offset=10)
k = cat(kcache, knew; dims=3)
v = cat(vcache, rand(Float32, 64, 2, 3, 4); dims=3)
y = scaled_dot_product_attention(q, k, v; is_causal=true, q_offset=10)Documents packed in one sequence (see scaled_dot_product_attention), whose positions start over at each document:
segments = [5, 12, 3]
positions = reduce(vcat, [1:n for n in segments]) # [1:5; 1:12; 1:3]
q = apply_rotary_embedding(rand(Float32, 64, 8, 20), cos_table, sin_table; positions)
k = apply_rotary_embedding(rand(Float32, 64, 2, 20), cos_table, sin_table; positions)
v = rand(Float32, 64, 2, 20)
y = scaled_dot_product_attention(q, k, v; is_causal=true, q_segments=segments)Softmax
Flux's logitcrossentropy uses NNlib.logsoftmax internally.
NNlib.softmax — Function
softmax(x; dims = 1)Softmax turns input array x into probability distributions that sum to 1 along the dimensions specified by dims. It is semantically equivalent to the following:
softmax(x; dims = 1) = exp.(x) ./ sum(exp.(x), dims = dims)with additional manipulations enhancing numerical stability. For Float16 and BFloat16 inputs, the sums are computed in Float32.
For a matrix input x it will by default (dims = 1) treat it as a batch of vectors, with each column independent. Keyword dims = 2 will instead treat rows independently, and so on.
See also logsoftmax.
Examples
julia> softmax([1, 2, 3])
3-element Vector{Float64}:
0.09003057317038046
0.24472847105479764
0.6652409557748218
julia> softmax([1 2 3; 2 2 2]) # dims=1
2×3 Matrix{Float64}:
0.268941 0.5 0.731059
0.731059 0.5 0.268941
julia> softmax([1 2 3; 2 2 2]; dims=2)
2×3 Matrix{Float64}:
0.0900306 0.244728 0.665241
0.333333 0.333333 0.333333Note that, when used with Flux.jl, softmax must not be passed to layers like Dense which accept an activation function. The activation is broadcasted over the result, thus applies to individual numbers. But softmax always needs to see the whole column.
julia> using Flux
julia> x = randn(Float32, 4, 4, 3, 13);
julia> model = Chain(Conv((4, 4), 3 => 8, tanh), Flux.flatten, Dense(8 => 7), softmax);
julia> model(x) |> size
(7, 13)
julia> Dense(4 => 7, softmax)(x)
ERROR: `softmax(x)` called with a number, but it expects an array. NNlib.softmax! — Function
softmax!(out, x; dims = 1)
softmax!(x; dims = 1)In-place version of softmax, writing the result into out (or x itself).
NNlib.∇softmax — Function
∇softmax(dy, y; dims = 1)Gradient of softmax with respect to its input, given the output y = softmax(x; dims) and the upstream cotangent dy. Does not depend on x.
NNlib.∇softmax! — Function
∇softmax!(dx, dy, y; dims = 1)In-place version of ∇softmax, writing the gradient into dx.
NNlib.logsoftmax — Function
logsoftmax(x; dims = 1)Computes the log of softmax in a more numerically stable way than directly taking log.(softmax(xs)). Commonly used in computing cross entropy loss.
It is semantically equivalent to the following:
logsoftmax(x; dims = 1) = x .- log.(sum(exp.(x), dims = dims))For Float16 and BFloat16 inputs, the sum is computed in Float32.
See also softmax.
NNlib.logsoftmax! — Function
logsoftmax!(out, x; dims = 1)
logsoftmax!(x; dims = 1)In-place version of logsoftmax, writing the result into out (or x itself).
NNlib.∇logsoftmax — Function
∇logsoftmax(dy, y; dims = 1)Gradient of logsoftmax with respect to its input, given the output y = logsoftmax(x; dims) and the upstream cotangent dy. Does not depend on x.
NNlib.∇logsoftmax! — Function
∇logsoftmax!(dx, dy, y; dims = 1)In-place version of ∇logsoftmax, writing the gradient into dx.
Pooling
Flux's AdaptiveMaxPool, AdaptiveMeanPool, GlobalMaxPool, GlobalMeanPool, MaxPool, MeanPool and lpnormpool use NNlib.PoolDims, NNlib.maxpool, NNlib.meanpool and NNlib.lpnormpool as their backend.
NNlib.meanpool supports complex inputs (the real and imaginary parts are pooled independently). This is handled in the generic code, so it works on every backend that pools real arrays — CPU, CUDA and AMDGPU. maxpool does not support complex inputs, since max is undefined for complex numbers.
NNlib.PoolDims — Type
PoolDims(x_size::NTuple{M}, k::Union{NTuple{<:Any, Int}, Int};
stride=k, padding=0, dilation=1, ceil_mode=false) where MDimensions for a "pooling" operation that can have an arbitrary input size, kernel size, stride, dilation, and channel count. Used to dispatch onto efficient implementations at compile-time.
With ceil_mode=true the output size is rounded up instead of down, as with PyTorch's ceil_mode=True: a last window that runs past the padded input is kept, as long as it starts inside the input or the low padding. Its positions past the padded input count neither as input nor as padding: max pooling skips them, and mean pooling leaves them out of the divisor even with count_include_pad=true. ceil_mode is dropped when it doesn't change the output size.
NNlib.maxpool — Function
maxpool(x, k::NTuple{N, Integer}; pad=0, stride=k, ceil_mode=false)Perform max pool operation with window size k on input tensor x.
Arguments:
xandk: Expectsndim(x) ∈ 3:5, and alwayslength(k) == ndim(x) - 2pad: Seepad_zerosfor details.stride: Either a tuple with the same length ask, or one integer for all directions. Default isk.ceil_mode: Whentrue, round the output size up instead of down, as PyTorch'sceil_mode=True: a last window that runs past the padded input is kept if it starts inside the input or the low padding. SeePoolDims.
NNlib.meanpool — Function
meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k, count_include_pad=true, ceil_mode=false)Perform mean pool operation with window size k on input tensor x.
Arguments:
xandk: Expectsndim(x) ∈ 3:5, and alwayslength(k) == ndim(x) - 2pad: Seepad_zerosfor details.stride: Either a tuple with the same length ask, or one integer for all directions. Default isk.count_include_pad: Whentrue(the default), padding (zero) elements are included in the averaging denominator, so each output element divides by the full window sizeprod(k). Whenfalse, only the in-bounds (non-padding) elements are counted, matching cuDNN'sCUDNN_POOLING_AVERAGE_COUNT_EXCLUDE_PADDINGand PyTorch'scount_include_pad=false. Has no effect whenpad == 0.ceil_mode: Whentrue, round the output size up instead of down, as PyTorch'sceil_mode=True: a last window that runs past the padded input is kept if it starts inside the input or the low padding. Its positions past the padded input are left out of the denominator also withcount_include_pad=true, as in PyTorch. SeePoolDims.
NNlib.lpnormpool — Function
lpnormpool(x, p::Real, k::NTuple{N, Integer}; pad=0, stride=k, ceil_mode=false)Perform Lp pool operation with value of the Lp norm p and window size k on input tensor x, also known as LPPool in pytorch. This pooling operator from Learned-Norm Pooling for Deep Feedforward and Recurrent Neural Networks.
Arguments:
xandk: Expectsndim(x) ∈ 3:5, and alwayslength(k) == ndim(x) - 2pis restricted to0 < p < Inf.pad: Seepad_zerosfor details.stride: Either a tuple with the same length ask, or one integer for all directions. Default isk.ceil_mode: Whentrue, round the output size up instead of down: a last window that runs past the padded input is kept if it starts inside the input or the low padding. SeePoolDims.
For all elements x in a size k window, lpnormpool computes (∑ᵢ xᵢ^p)^(1 / p) as an element of the output.
Thus lpnormpool(x, 1, k) ./ prod(k) ≈ meanpool(x, k) and lpnormpool(x, 2, k).^2 ./ prod(k) ≈ meanpool(x.^2, k). With ceil_mode=true, the sum of a last window that runs past the input is over its elements in the input; PyTorch's LPPool instead scales it by prod(k) over their number.
Padding
NNlib.pad_reflect — Function
pad_reflect(x, pad::Tuple; [dims])
pad_reflect(x, pad::Int; [dims])Pad the array x reflecting its values across the border.
pad can a tuple of integers (l1, r1, ..., ln, rn) of some length 2n that specifies the left and right padding size for each of the dimensions in dims. If dims is not given, it defaults to the first n dimensions.
If pad is an integer, it is applied on both sides on every dimension in dims. In this case, dims defaults to the first ndims(x)-2 dimensions (i.e. excludes the channel and batch dimension).
See also pad_repeat, pad_symmetric, pad_circular, and pad_constant.
julia> r = reshape(1:9, 3, 3)
3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64:
1 4 7
2 5 8
3 6 9
julia> pad_reflect(r, (1,2,1,2))
6×6 Matrix{Int64}:
5 2 5 8 5 2
4 1 4 7 4 1
5 2 5 8 5 2
6 3 6 9 6 3
5 2 5 8 5 2
4 1 4 7 4 1NNlib.pad_symmetric — Function
pad_symmetric(x, pad::Tuple; [dims])
pad_symmetric(x, pad::Int; [dims])Pad the array x reflecting its values symmetrically across the border, i.e. the border values of x are present in the padding values, in contrast to pad_reflect.
pad can a tuple of integers (l1, r1, ..., ln, rn) of some length 2n that specifies the left and right padding size for each of the dimensions in dims. If dims is not given, it defaults to the first n dimensions.
If pad is an integer, it is applied on both sides on every dimension in dims. In this case, dims defaults to the first ndims(x)-2 dimensions (i.e. excludes the channel and batch dimension).
See also pad_repeat, pad_reflect, pad_circular, and pad_constant.
julia> r = reshape(1:9, 3, 3)
3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64:
1 4 7
2 5 8
3 6 9
julia> pad_symmetric(r, (1,2,1,2))
6×6 Matrix{Int64}:
1 1 4 7 7 4
1 1 4 7 7 4
2 2 5 8 8 5
3 3 6 9 9 6
3 3 6 9 9 6
2 2 5 8 8 5NNlib.pad_circular — Function
pad_circular(x, pad::Tuple; [dims])
pad_circular(x, pad::Int; [dims])Pad the array x "circularly" across the border by wrapping around values from the opposite side of x.
pad can a tuple of integers (l1, r1, ..., ln, rn) of some length 2n that specifies the left and right padding size for each of the dimensions in dims. If dims is not given, it defaults to the first n dimensions.
If pad is an integer, it is applied on both sides on every dimension in dims. In this case, dims defaults to the first ndims(x)-2 dimensions (i.e. excludes the channel and batch dimension).
The pad length on either side in any dimension must not exceed the size of x in that dimension, i.e. pad_circular is not able to create abitrary sized tilings of x.
See also pad_repeat, pad_reflect, pad_symmetric, and pad_constant.
julia> r = reshape(1:9, 3, 3)
3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64:
1 4 7
2 5 8
3 6 9
julia> pad_circular(r, (1,2,1,2))
6×6 Matrix{Int64}:
9 3 6 9 3 6
7 1 4 7 1 4
8 2 5 8 2 5
9 3 6 9 3 6
7 1 4 7 1 4
8 2 5 8 2 5NNlib.pad_repeat — Function
pad_repeat(x, pad::Tuple; [dims])
pad_repeat(x, pad::Int; [dims])Pad the array x repeating the values on the border.
pad can a tuple of integers (l1, r1, ..., ln, rn) of some length 2n that specifies the left and right padding size for each of the dimensions in dims. If dims is not given, it defaults to the first n dimensions.
If pad is an integer, it is applied on both sides on every dimension in dims. In this case, dims defaults to the first ndims(x)-2 dimensions (i.e. excludes the channel and batch dimension).
See also pad_reflect, pad_symmetric, pad_circular, and pad_constant.
julia> r = reshape(1:9, 3, 3)
3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64:
1 4 7
2 5 8
3 6 9
julia> pad_repeat(r, (1,2,3,4))
6×10 Matrix{Int64}:
1 1 1 1 4 7 7 7 7 7
1 1 1 1 4 7 7 7 7 7
2 2 2 2 5 8 8 8 8 8
3 3 3 3 6 9 9 9 9 9
3 3 3 3 6 9 9 9 9 9
3 3 3 3 6 9 9 9 9 9NNlib.pad_constant — Function
pad_constant(x, pad::Tuple, val = 0; [dims = :])
pad_constant(x, pad::Int, val = 0; [dims = :])Pad the array x with the constant value val.
pad can be a tuple of integers. If it is of some length 2 * length(dims) that specifies the left and right padding size for each of the dimensions in dims as (l1, r1, ..., ln, rn). If supplied with a tuple of length length(dims) instead, it applies symmetric padding. If dims is not given, it defaults to all dimensions.
For integer pad input, it is applied on both sides on every dimension in dims.
See also pad_zeros, pad_repeat, pad_reflect, pad_symmetric, and pad_circular.
julia> r = reshape(1:4, 2, 2)
2×2 reshape(::UnitRange{Int64}, 2, 2) with eltype Int64:
1 3
2 4
julia> pad_constant(r, (1, 2, 3, 4), 8)
5×9 Matrix{Int64}:
8 8 8 8 8 8 8 8 8
8 8 8 1 3 8 8 8 8
8 8 8 2 4 8 8 8 8
8 8 8 8 8 8 8 8 8
8 8 8 8 8 8 8 8 8
julia> pad_constant(r, 1, 8)
4×4 Matrix{Int64}:
8 8 8 8
8 1 3 8
8 2 4 8
8 8 8 8
julia> r = reshape(1:27, 3, 3, 3)
3×3×3 reshape(::UnitRange{Int64}, 3, 3, 3) with eltype Int64:
[:, :, 1] =
1 4 7
2 5 8
3 6 9
[:, :, 2] =
10 13 16
11 14 17
12 15 18
[:, :, 3] =
19 22 25
20 23 26
21 24 27
julia> pad_constant(r, (2,1), dims = 1) # assymetric padding
6×3×3 Array{Int64, 3}:
[:, :, 1] =
0 0 0
0 0 0
1 4 7
2 5 8
3 6 9
0 0 0
[:, :, 2] =
0 0 0
0 0 0
10 13 16
11 14 17
12 15 18
0 0 0
[:, :, 3] =
0 0 0
0 0 0
19 22 25
20 23 26
21 24 27
0 0 0
julia> pad_constant(r, (2,1, 3), dims = (1,2)) # padding must always be either the same length as dims, or double it
ERROR: ArgumentError: Could not parse padding (2, 1, 3) and dims (1, 2)
Stacktrace:
[...]NNlib.pad_zeros — Function
pad_zeros(x, pad::Tuple; [dims])
pad_zeros(x, pad::Int; [dims])Pad the array x with zeros. Equivalent to pad_constant with the constant equal to 0.
Convolution
Flux's Conv and CrossCor layers use NNlib.DenseConvDims and NNlib.conv internally.
NNlib.conv supports complex datatypes on CPU and CUDA devices.
Therefore, for every regular convolution (flipped=false), the kernel is flipped before calculation. For better performance, use cross-correlation (flipped=true) and manually flip the kernel before the NNlib.conv call. Flux handles this automatically; this is only required for direct calls.
NNlib.conv — Function
conv(x, w; stride = 1, pad = 0, dilation = 1, flipped = false, groups = 1)Apply convolution filter w to input x. x and w are 3d/4d/5d tensors in 1d/2d/3d convolutions respectively. They may have real or complex element types.
The dimensions are arranged as follows, with the spatial dimensions first:
x:(spatial..., channels_in, batch)w:(spatial..., channels_in ÷ groups, channels_out)- output:
(spatial_out..., channels_out, batch)
The output spatial size for each dimension is (spatial + 2*pad - (dilation*(spatial_w - 1) + 1)) ÷ stride + 1.
Keyword arguments
stride: stride of the filter (Intor tuple of oneIntper spatial dimension).pad: zero padding. AnIntor tuple applies symmetric padding; a tuple of2*ndimsentries gives(left, right)padding per spatial dimension.dilation: spacing between filter elements (atrous convolution).flipped: iffalse(default), perform true mathematical convolution; iftrue, perform cross-correlation (the usual deep-learning convention).groups: number of groups for grouped convolution.
Examples
julia> x = rand(Float32, 28, 28, 3, 16); # (width, height, channels_in, batch)
julia> w = rand(Float32, 5, 5, 3, 8); # (width, height, channels_in, channels_out)
julia> size(conv(x, w)) # 28 - 5 + 1 = 24
(24, 24, 8, 16)
julia> size(conv(x, w; stride=2, pad=1))
(13, 13, 8, 16)NNlib.ConvDims — Type
ConvDimsType system-level information about convolution dimensions. Critical for things like im2col!() to generate efficient code, and helpful to reduce the number of kwargs getting passed around.
NNlib.depthwiseconv — Function
depthwiseconv(x, w; stride=1, pad=0, dilation=1, flipped=false)Depthwise convolution operation with filter w on input x. x and w are 3d/4d/5d tensors in 1d/2d/3d convolutions respectively. The filter w has shape (spatial..., channel_multiplier, channels_in).
This is equivalent to a grouped conv with groups = channels_in, and is implemented in terms of it.
NNlib.DepthwiseConvDims — Type
DepthwiseConvDimsConcrete subclass of ConvDims for a depthwise convolution. Differs primarily due to characterization by C_in, C_mult, rather than C_in, C_out. Useful to be separate from DenseConvDims primarily for channel calculation differences.
NNlib.DenseConvDims — Type
DenseConvDimsConcrete subclass of ConvDims for a normal, dense, conv2d/conv3d.
NNlib.unfold — Function
unfold(x, kernel_size; stride = 1, pad = 0, dilation = 0, flipped = true)Places sliding windows of x into a container tensor of size (num_windows, window_size, batchsize). The window size is determined by the prod(spatial dims of kernel)*input_channels. The number of sliding windows will match those of convolution (conv) with the same kernel_size and arguments. Note that by default conv flips the spatial dimensions of its kernel (default flipped=false), whereas unfold does not (default flipped=true). Uses NNlib.im2col! as backend.
See also fold, the adjoint/transpose operator and a potential inverse of unfold.
Example
The below example demonstrates that unfold uses the same sliding windows as conv. In general batched_mul + unfold should not be used to achieve convolution.
julia> x = reshape([100 2 3 40 5 6 700], 7, 1, 1); # 1D data, 1 channel, batch of 1
julia> w = reshape([1 0 -1], 3, 1, 1); # 1D conv kernel of length 3
julia> kws = (pad=1, stride=2, flipped=true); # use same args for conv and unfold
julia> z = NNlib.unfold(x, size(w); kws...)
4×3×1 Array{Int64, 3}:
[:, :, 1] =
0 100 2
2 3 40
40 5 6
6 700 0
julia> y1 = conv(x, w; kws...)
4×1×1 Array{Int64, 3}:
[:, :, 1] =
-2
-38
34
6
julia> y2 = z ⊠ w # ⊠ (\boxtimes) is NNlib.batched_mul
4×1×1 Array{Int64, 3}:
[:, :, 1] =
-2
-38
34
6NNlib.fold — Function
fold(y, output_size, kernel_size; stride = 1, pad = 0, dilation = 0, flipped = true)The adjoint/transpose operator of unfold. It accumulates sliding windows from the output of unfold into a container tensor of size output_size. An inverse to unfold may be obtained (in some cases) by using fold and accounting for scaling issues with a divisor (see example). Uses NNlib.col2im! as backend.
See also unfold.
Example
julia> x = reshape([100 2 3 40 5 6 700], 7, 1, 1); # 1D data, 1 channel, batch of 1
julia> y = NNlib.unfold(x, (3,1,1)) # sliding window of size 3
5×3×1 Array{Int64, 3}:
[:, :, 1] =
100 2 3
2 3 40
3 40 5
40 5 6
5 6 700
julia> z = NNlib.fold(y, size(x), (3,1,1)) # sum of contributions in y. 100 appears once, 40 three times
7×1×1 Array{Int64, 3}:
[:, :, 1] =
100
4
9
120
15
12
700
julia> divisor = NNlib.fold(NNlib.unfold(ones(size(x)...), (3,1,1)), size(x), (3,1,1))
7×1×1 Array{Float64, 3}:
[:, :, 1] =
1.0
2.0
3.0
3.0
3.0
2.0
1.0
julia> z ./ divisor
7×1×1 Array{Float64, 3}:
[:, :, 1] =
100.0
2.0
3.0
40.0
5.0
6.0
700.0In general, an inverse to unfold does not exist if divisor contains zeros.
Upsampling
Flux's Upsample layer uses NNlib.upsample_nearest, NNlib.upsample_bilinear, and NNlib.upsample_trilinear as its backend. Additionally, Flux's PixelShuffle layer uses NNlib.pixel_shuffle as its backend.
NNlib.upsample_nearest — Function
upsample_nearest(x, scale::NTuple{S,Int})
upsample_nearest(x; size::NTuple{S,Int})Upsamples the array x by integer multiples along the first S dimensions. Subsequent dimensions of x are not altered.
Either the scale factors or the final output size can be specified.
See also upsample_bilinear, for two dimensions of an N=4 array.
Example
julia> upsample_nearest([1 2 3; 4 5 6], (2, 3))
4×9 Matrix{Int64}:
1 1 1 2 2 2 3 3 3
1 1 1 2 2 2 3 3 3
4 4 4 5 5 5 6 6 6
4 4 4 5 5 5 6 6 6
julia> ans == upsample_nearest([1 2 3; 4 5 6]; size=(4, 9)) # equivalent
true
julia> upsample_nearest([1 2 3; 4 5 6], (2,))
4×3 Matrix{Int64}:
1 2 3
1 2 3
4 5 6
4 5 6
julia> ans == upsample_nearest([1 2 3; 4 5 6], size=(4,))
trueNNlib.∇upsample_nearest — Function
∇upsample_nearest(Δ::AbstractArray{T,3}, scales::NTuple{S, <:Integer}) where TArguments
Δ: Incoming gradient array, backpropagated from downstream layersscales: scales by which the image was upsampled in the first place
Outputs
dx: Downsampled version ofΔ
NNlib.upsample_linear — Function
upsample_linear(x::AbstractArray{T,3}, scale::Real; align_corners::Bool = true)
upsample_linear(x::AbstractArray{T,3}; size::Integer, align_corners::Bool = true)Upsamples the first dimension of the array x by the upsample provided scale, using linear interpolation. As an alternative to using scale, the resulting array size can be directly specified with a keyword argument.
The size of the output is equal to (scale*S1, S2, S3), where S1, S2, S3 = size(x).
NNlib.∇upsample_linear — Function
∇upsample_linear(Δ::AbstractArray{T,3}; size::Integer, align_corners::Bool = true) where TArguments
Δ: Incoming gradient array, backpropagated from downstream layerssize: Size of the image upsampled in the first place
Outputs
dx: Downsampled version ofΔ
NNlib.upsample_bilinear — Function
upsample_bilinear(x::AbstractArray{T,4}, scale::NTuple{2,Real}; align_corners::Bool = true, antialias::Bool = false)
upsample_bilinear(x::AbstractArray{T,4}; size::NTuple{2,Integer}, align_corners::Bool = true, antialias::Bool = false)Upsamples the first 2 dimensions of the array x by the upsample factors stored in scale, using bilinear interpolation. As an alternative to using scale, the resulting image size can be directly specified with a keyword argument.
The size of the output is equal to (scale[1]*S1, scale[2]*S2, S3, S4), where S1, S2, S3, S4 = size(x).
With antialias = true, the output can also be smaller than the input: as in PyTorch's F.interpolate(x, size, mode="bilinear", antialias=True) and PIL, the triangle filter is then stretched by the downsampling factor and renormalised at the borders. Pass align_corners = false to match them. See also upsample_bicubic.
Examples
julia> x = reshape(Float32[1 2 3; 4 5 6], (2,3,1,1))
2×3×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
1.0 2.0 3.0
4.0 5.0 6.0
julia> upsample_bilinear(x, (2, 3))
4×9×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
1.0 1.25 1.5 1.75 2.0 2.25 2.5 2.75 3.0
2.0 2.25 2.5 2.75 3.0 3.25 3.5 3.75 4.0
3.0 3.25 3.5 3.75 4.0 4.25 4.5 4.75 5.0
4.0 4.25 4.5 4.75 5.0 5.25 5.5 5.75 6.0
julia> ans == upsample_bilinear(x; size=(4, 9)) # specify ouput size instead
true
julia> upsample_bilinear(x, (2.5, 3.5)) # non-integer scaling factors are allowed
5×10×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
1.0 1.22222 1.44444 1.66667 1.88889 … 2.33333 2.55556 2.77778 3.0
1.75 1.97222 2.19444 2.41667 2.63889 3.08333 3.30556 3.52778 3.75
2.5 2.72222 2.94444 3.16667 3.38889 3.83333 4.05556 4.27778 4.5
3.25 3.47222 3.69444 3.91667 4.13889 4.58333 4.80556 5.02778 5.25
4.0 4.22222 4.44444 4.66667 4.88889 5.33333 5.55556 5.77778 6.0NNlib.∇upsample_bilinear — Function
∇upsample_bilinear(Δ::AbstractArray{T,4}; size::NTuple{2,Integer}, align_corners::Bool = true, antialias::Bool = false) where TArguments
Δ: Incoming gradient array, backpropagated from downstream layerssize: Lateral (W,H) size of the image upsampled in the first place
Outputs
dx: Downsampled version ofΔ
NNlib.upsample_bicubic — Function
upsample_bicubic(x::AbstractArray{T,4}, scale::NTuple{2,Real}; align_corners::Bool = false, antialias::Bool = false)
upsample_bicubic(x::AbstractArray{T,4}; size::NTuple{2,Integer}, align_corners::Bool = false, antialias::Bool = false)Resizes the first 2 dimensions of the array x with bicubic interpolation, to the image size given as a keyword argument, or to floor.(Int, scale .* size(x)[1:2]). The output can be smaller than the input.
This is PyTorch's F.interpolate(x, size, mode="bicubic", align_corners=..., antialias=...) on the same data in the NCHW layout. Without antialias, each output pixel is a weighted sum of 4×4 input pixels (the cubic convolution kernel with a = -0.75, as in OpenCV), and the image is extended by repeating its border pixels. With antialias = true, the kernel is the one of PIL and torchvision's Resize (a = -0.5): when downsampling, it is stretched by the downsampling factor, so that each output pixel averages all the input pixels it covers, and at the borders the weights are renormalised over the pixels inside the image.
Unlike upsample_bilinear, align_corners defaults to false, as in PyTorch.
Examples
julia> x = reshape(Float32[1 2; 3 4], (2, 2, 1, 1));
julia> upsample_bicubic(x, (2, 2)) # overshoots the range of `x`
4×4×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
0.683594 1.01562 1.5625 1.89453
1.34766 1.67969 2.22656 2.55859
2.44141 2.77344 3.32031 3.65234
3.10547 3.4375 3.98438 4.31641
julia> ans == upsample_bicubic(x; size = (4, 4))
true
julia> x = reshape(Float32.(1:16), (4, 4, 1, 1));
julia> upsample_bicubic(x; size = (2, 2), antialias = true)
2×2×1×1 Array{Float32, 4}:
[:, :, 1, 1] =
3.93388 11.2397
5.76033 13.0661NNlib.∇upsample_bicubic — Function
∇upsample_bicubic(Δ::AbstractArray{T,4}; size::NTuple{2,Integer}, align_corners::Bool = false, antialias::Bool = false)Arguments
Δ: Incoming gradient array, backpropagated from downstream layerssize: Lateral (W,H) size of the image resized in the first place
Outputs
dx: Gradient with respect to the input ofupsample_bicubic
NNlib.upsample_trilinear — Function
upsample_trilinear(x::AbstractArray{T,5}, scale::NTuple{3,Real}; align_corners::Bool = true)
upsample_trilinear(x::AbstractArray{T,5}; size::NTuple{3,Integer}, align_corners::Bool = true)Upsamples the first 3 dimensions of the array x by the upsample factors stored in scale, using trilinear interpolation. As an alternative to using scale, the resulting image size can be directly specified with a keyword argument.
The size of the output is equal to (scale[1]*S1, scale[2]*S2, scale[3]*S3, S4, S5), where S1, S2, S3, S4, S5 = size(x).
Examples
upsample_trilinear(x, (2, 3, 4))
upsample_trilinear(x; size=(4, 9, 11)) # specify ouput size instead
upsample_trilinear(x, (2.5, 3.5, pi)) # non-integer scaling factors are allowedNNlib.∇upsample_trilinear — Function
∇upsample_trilinear(Δ::AbstractArray{T,5}; size::NTuple{3,Integer}, align_corners::Bool = true) where TArguments
Δ: Incoming gradient array, backpropagated from downstream layerssize: Lateral size & depth (W,H,D) of the image upsampled in the first place
Outputs
dx: Downsampled version ofΔ
NNlib.pixel_shuffle — Function
pixel_shuffle(x, r::Integer)Pixel shuffling operation, upscaling by a factor r.
For 4-arrays representing N images, the operation converts input size(x) == (W, H, r^2*C, N) to output of size (r*W, r*H, C, N). For D-dimensional data, it expects ndims(x) == D+2 with channel and batch dimensions, and divides the number of channels by r^D.
Used in super-resolution networks to upsample towards high resolution features. Reference: Shi et. al., "Real-Time Single Image and Video Super-Resolution ...", CVPR 2016, https://arxiv.org/abs/1609.05158
Examples
julia> x = [10i + j + channel/10 for i in 1:2, j in 1:3, channel in 1:4, batch in 1:1]
2×3×4×1 Array{Float64, 4}:
[:, :, 1, 1] =
11.1 12.1 13.1
21.1 22.1 23.1
[:, :, 2, 1] =
11.2 12.2 13.2
21.2 22.2 23.2
[:, :, 3, 1] =
11.3 12.3 13.3
21.3 22.3 23.3
[:, :, 4, 1] =
11.4 12.4 13.4
21.4 22.4 23.4
julia> pixel_shuffle(x, 2) # 4 channels used up as 2x upscaling of image dimensions
4×6×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
11.1 11.3 12.1 12.3 13.1 13.3
11.2 11.4 12.2 12.4 13.2 13.4
21.1 21.3 22.1 22.3 23.1 23.3
21.2 21.4 22.2 22.4 23.2 23.4
julia> y = [i + channel/10 for i in 1:3, channel in 1:6, batch in 1:1]
3×6×1 Array{Float64, 3}:
[:, :, 1] =
1.1 1.2 1.3 1.4 1.5 1.6
2.1 2.2 2.3 2.4 2.5 2.6
3.1 3.2 3.3 3.4 3.5 3.6
julia> pixel_shuffle(y, 2) # 1D image, with 6 channels reduced to 3
6×3×1 Array{Float64, 3}:
[:, :, 1] =
1.1 1.3 1.5
1.2 1.4 1.6
2.1 2.3 2.5
2.2 2.4 2.6
3.1 3.3 3.5
3.2 3.4 3.6Rotation
Rotate images in the first two dimensions of an array.
NNlib.imrotate — Function
imrotate(arr::AbstractArray{T, 4}, θ; method=:bilinear, rotation_center=size(arr) .÷ 2 .+ 1)Rotates an array in the first two dimensions around the center pixel rotation_center. The default value of rotation_center is defined such that there is a integer center pixel for even and odd sized arrays which it is rotated around. For an even sized array of size (4,4) this would be (3,3), for an odd array of size (3,3) this would be (2,2) However, rotation_center can be also non-integer numbers if specified.
The angle θ is interpreted in radians.
The adjoint is defined with ChainRulesCore.jl. This method also runs with CUDA (and in principle all KernelAbstractions.jl supported backends).
Keywords
method=:bilinearfor bilinear interpolation ormethod=:nearestfor nearest neighbourrotation_center=size(arr) .÷ 2 .+ 1means there is a real center pixel around it is rotated.
Examples
julia> arr = zeros((4,4,1,1)); arr[2,2,1,1] = 1;
julia> arr
4×4×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
0.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0
julia> NNlib.imrotate(arr, deg2rad(90)) # rotation around (3,3)
4×4×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
0.0 0.0 0.0 0.0
0.0 0.0 0.0 1.0
0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0
julia> NNlib.imrotate(arr, deg2rad(90), rotation_center=(2,2))
4×4×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
0.0 0.0 0.0 0.0
0.0 1.0 0.0 0.0
0.0 0.0 0.0 0.0
0.0 0.0 0.0 0.0
julia> arr = zeros((3,3,1,1)); arr[1,2,1,1] = 1
1
julia> arr
3×3×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
0.0 1.0 0.0
0.0 0.0 0.0
0.0 0.0 0.0
julia> NNlib.imrotate(arr, deg2rad(45))
3×3×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
0.0 0.207107 0.0
0.0 0.0 0.207107
0.0 0.0 0.0
julia> NNlib.imrotate(arr, deg2rad(45), method=:nearest)
3×3×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
0.0 0.0 1.0
0.0 0.0 0.0
0.0 0.0 0.0NNlib.∇imrotate — Function
∇imrotate(dy, arr::AbstractArray{T, 4}, θ; method=:bilinear,
rotation_center=size(arr) .÷ 2 .+ 1)Adjoint for imrotate. Gradient only with respect to arr and not θ.
Arguments
dy: input gradientarr: Input from primal computationθ: rotation angle in radiansmethod=:bilinearormethod=:nearestrotation_center=size(arr) .÷ 2 .+ 1rotates around a real center pixel for even and odd sized arrays
Batched Operations
Flux's Bilinear layer uses NNlib.batched_mul internally.
NNlib.batched_mul — Function
batched_mul(A, B) -> C
A ⊠ B # \boxtimesBatched matrix multiplication. Result has C[:,:,k...] == A[:,:,k...] * B[:,:,k...] where k... represent any indices in the last dimensions.
If ndims(A) == ndims(B) == 3 and size(B,3) == 1 then instead C[:,:,k] == A[:,:,k] * B[:,:,1], and similarly for A.
To transpose each matrix, apply batched_transpose to the array, or batched_adjoint for conjugate-transpose:
julia> A, B = randn(2,5,17), randn(5,9,17);
julia> A ⊠ B |> size
(2, 9, 17)
julia> batched_adjoint(A) |> size
(5, 2, 17)
julia> batched_mul(A, batched_adjoint(randn(9,5,17))) |> size
(2, 9, 17)
julia> A ⊠ randn(5,9,1) |> size
(2, 9, 17)
julia> batched_transpose(A) == PermutedDimsArray(A, (2,1,3))
trueThe equivalent PermutedDimsArray may be used in place of batched_transpose. Other permutations are also handled by BLAS, provided that the batch index k is not the first dimension of the underlying array. Thus PermutedDimsArray(::Array, (1,3,2)) and PermutedDimsArray(::Array, (3,1,2)) are fine.
However, A = PermutedDimsArray(::Array, (3,2,1)) is not acceptable to BLAS, since the batch dimension is the contiguous one: stride(A,3) == 1. This will be copied, as doing so is faster than batched_mul_generic!.
Both this copy and batched_mul_generic! produce @debug messages, and setting for instance ENV["JULIA_DEBUG"] = NNlib will display them.
batched_mul(A::Array{T,3}, B::Matrix)
batched_mul(A::Matrix, B::Array{T,3})
A ⊠ BThis is always matrix-matrix multiplication, but either A or B may lack a batch index.
When
Bis a matrix, result hasC[:,:,k] == A[:,:,k] * B[:,:]for allk.When
Ais a matrix, thenC[:,:,k] == A[:,:] * B[:,:,k]. This can also be done by reshaping and calling*, for instanceA ⊡ Busing TensorCore.jl, but is implemented here usingbatched_gemminstead ofgemm.
julia> randn(16,8,32) ⊠ randn(8,4) |> size
(16, 4, 32)
julia> randn(16,8,32) ⊠ randn(8,4,1) |> size # equivalent
(16, 4, 32)
julia> randn(16,8) ⊠ randn(8,4,32) |> size
(16, 4, 32)See also batched_vec to regard B as a batch of vectors, A[:,:,k] * B[:,k].
NNlib.batched_mul! — Function
batched_mul!(C, A, B) -> C
batched_mul!(C, A, B, α=1, β=0)In-place batched matrix multiplication, equivalent to mul!(C[:,:,k], A[:,:,k], B[:,:,k], α, β) for all k. If size(B,3) == 1 then every batch uses B[:,:,1] instead.
This will call batched_gemm! whenever possible. For real arrays this means that, for X ∈ [A,B,C], either stride(X,1)==1 or stride(X,2)==1, the latter may be caused by batched_transpose or by for instance PermutedDimsArray(::Array, (3,1,2)). Unlike batched_mul this will never make a copy.
For complex arrays, the wrapper made by batched_adjoint must be outermost to be seen. In this case the strided accepted by BLAS are more restricted, if stride(C,1)==1 then only stride(AorB::BatchedAdjoint,2) == 1 is accepted.
NNlib.batched_adjoint — Function
batched_transpose(A::AbstractArray{T,3})
batched_adjoint(A)Equivalent to applying transpose or adjoint to each matrix A[:,:,k].
These exist to control how batched_mul behaves, as it operates on such matrix slices of an array with ndims(A)==3.
PermutedDimsArray(A, (2,1,3)) is equivalent to batched_transpose(A), and is also understood by batched_mul (and more widely supported elsewhere).
BatchedTranspose{T, S} <: AbstractBatchedMatrix{T, 3}
BatchedAdjoint{T, S}Lazy wrappers analogous to Transpose and Adjoint, returned by batched_transpose etc.
NNlib.batched_transpose — Function
batched_transpose(A::AbstractArray{T,3})
batched_adjoint(A)Equivalent to applying transpose or adjoint to each matrix A[:,:,k].
These exist to control how batched_mul behaves, as it operates on such matrix slices of an array with ndims(A)==3.
PermutedDimsArray(A, (2,1,3)) is equivalent to batched_transpose(A), and is also understood by batched_mul (and more widely supported elsewhere).
BatchedTranspose{T, S} <: AbstractBatchedMatrix{T, 3}
BatchedAdjoint{T, S}Lazy wrappers analogous to Transpose and Adjoint, returned by batched_transpose etc.
NNlib.batched_vec — Function
batched_vec(A::AbstractArray{T,3}, B::AbstractMatrix)
batched_vec(A::AbstractArray{T,3}, b::AbstractVector)
batched_vec(A::AbstractArray, B::AbstractArray)Batched matrix-vector multiplication. For the 3D case: the result has C[:,:,k] == A[:,:,k] * B[:,k] for all k, or else C[:,:,k] == A[:,:,k] * b for b::Vector.
For the general N-D case where ndims(A) == ndims(B) + 1: the result has C[:,k...] == A[:,:,k...] * B[:,k...] for all batch indices k.... The batch dimensions must match: size(A)[3:end] == size(B)[2:end].
With the same argument types, batched_mul(A, B) would regard B as a fixed matrix, not a batch of vectors. Both reshape and then call batched_mul(::Array{T,3}, ::Array{T,3}).
julia> A, B, b = randn(16,8,32), randn(8,32), randn(8);
julia> batched_vec(A,B) |> size
(16, 32)
julia> batched_vec(A,b) |> size
(16, 32)
julia> A4d, B3d = randn(16,8,10,32), randn(8,10,32); # 4D and 3D arrays
julia> batched_vec(A4d, B3d) |> size
(16, 10, 32)NNlib.grouped_mul — Function
grouped_mul(W::AbstractArray{<:Any,3}, X::AbstractMatrix, offsets) -> Y
grouped_mul(A::AbstractMatrix, B::AbstractMatrix, offsets) -> CMatrix multiplication over groups of different sizes. offsets is an integer vector of nondecreasing cumulative group ends, starting from an implicit zero. Repeated offsets represent empty groups.
For W of size (n, k, G) and X of size (k, total), return Y of size (n, total), with Y[:, r] = W[:, :, g] * X[:, r] for each group g, where r = (g == 1 ? 0 : offsets[g-1])+1:offsets[g]. length(offsets) must equal G. W may be a batched_transpose or batched_adjoint.
For A of size (n, total) and B of size (total, k), return C of size (n, k, length(offsets)), with C[:, :, g] = A[:, r] * B[r, :].
Entries past the last offset are ignored. The corresponding columns of Y, and the matrices of C for empty groups, are zero. Both forms support differentiation.
On CUDA 6.5+, supported strided arrays use a single cuBLAS grouped GEMM, with Float32 accumulation for Float16 and BFloat16. Device offsets are copied to the host once per call, since cuBLAS requires group sizes on the host. Other arrays use a loop of mul! calls, parallelized over groups on multithreaded CPUs.
julia> W = cat(ones(2, 3), fill(2.0, 2, 3); dims=3);
julia> grouped_mul(W, ones(3, 4), [1, 3])
2×4 Matrix{Float64}:
3.0 6.0 6.0 0.0
3.0 6.0 6.0 0.0
julia> size(grouped_mul(ones(2, 4), ones(4, 3), [1, 3]))
(2, 3, 2)NNlib.grouped_mul! — Function
grouped_mul!(C, A, B, offsets, α=1, β=0) -> CCompute C = α * grouped_mul(A, B, offsets) + β * C in place. The destination must have the shape returned by grouped_mul and must not alias either input. Entries outside the groups are scaled by β. α and β are converted to the element type of C.
Gather and Scatter
Flux's Embedding layer uses NNlib.gather as its backend.
NNlib.gather — Function
NNlib.gather(src, idx) -> dstReverse operation of scatter. Gathers data from source src and writes it in a destination dst according to the index array idx. For each k in CartesianIndices(idx), assign values to dst according to
dst[:, ... , k] .= src[:, ... , idx[k]...]Notice that if idx is a vector containing integers and src is a matrix, previous expression simplifies to
dst[:, k] .= src[:, idx[k]]and k will run over 1:length(idx).
The elements of idx can be integers or integer tuples and may be repeated. A single src column can end up being copied into zero, one, or multiple dst columns.
See gather! for an in-place version.
Examples
julia> NNlib.gather([1,20,300,4000], [2,4,2])
3-element Vector{Int64}:
20
4000
20
julia> NNlib.gather([1 2 3; 4 5 6], [1,3,1,3,1])
2×5 Matrix{Int64}:
1 3 1 3 1
4 6 4 6 4gather(src, IJK...)Convert the tuple of integer vectors IJK to a tuple of CartesianIndex and call gather on it: gather(src, CartesianIndex.(IJK...)).
Examples
julia> src = reshape([1:15;], 3, 5)
3×5 Matrix{Int64}:
1 4 7 10 13
2 5 8 11 14
3 6 9 12 15
julia> NNlib.gather(src, [1, 2], [2, 4])
2-element Vector{Int64}:
4
11NNlib.gather! — Function
NNlib.gather!(dst, src, idx)Reverse operation of scatter!. Gathers data from source src and writes it in destination dst according to the index array idx. For each k in CartesianIndices(idx), assign values to dst according to
dst[:, ... , k] .= src[:, ... , idx[k]...]Notice that if idx is a vector containing integers, and both dst and src are matrices, previous expression simplifies to
dst[:, k] .= src[:, idx[k]]and k will run over 1:length(idx).
The elements of idx can be integers or integer tuples and may be repeated. A single src column can end up being copied into zero, one, or multiple dst columns.
See gather for an allocating version.
gather!(dst, src, IJK...)Convert the tuple of integer vectors IJK to a tuple of CartesianIndex and call gather! on it: gather!(dst, src, CartesianIndex.(IJK...)).
Examples
julia> src = reshape([1:15;], 3, 5)
3×5 Matrix{Int64}:
1 4 7 10 13
2 5 8 11 14
3 6 9 12 15
julia> dst = zeros(Int, 2);
julia> NNlib.gather!(dst, src, [1, 2], [2, 4])
2-element Vector{Int64}:
4
11NNlib.scatter — Function
NNlib.scatter(op, src, idx; [init, dstsize])Scatter operation allocating a destination array dst and calling scatter!(op, dst, src, idx) on it.
If keyword
initis provided, it is used to initialize the content ofdst. Otherwise, the init values is inferred from the reduction operatoropfor some common operators (e.g.init = 0forop = +).If
dstsizeis provided, it will be used to define the size of destination array, otherwise it will be inferred bysrcandidx.
See scatter! for full details on how idx works.
Examples
julia> NNlib.scatter(+, [10,100,1000], [3,1,2])
3-element Vector{Int64}:
100
1000
10
julia> NNlib.scatter(+, [1 2 3 4; 5 6 7 8], [2,1,1,5])
2×5 Matrix{Int64}:
5 1 0 0 4
13 5 0 0 8
julia> NNlib.scatter(*, [10,200,3000], [1,4,2]; init = 10, dstsize = 6)
6-element Vector{Int64}:
100
30000
10
2000
10
10NNlib.scatter! — Function
NNlib.scatter!(op, dst, src, idx)Scatter operation, which writes data in src into dst at locations idx. A binary reduction operator op is applied during the scatter. For each index k in idx, accumulates values in dst according to
dst[:, ..., idx[k]...] = (op).(dst[:, ..., idx[k]...], src[:, ..., k...])Arguments
op: Operations to be applied ondstandsrc, e.g.+,-,*,/,max,minandmean.dst: The destination forsrcto aggregate to. This argument will be mutated.src: The source data for aggregating.idx: The mapping for aggregation from source (index) to destination (value). Theidxarray can contain either integers or tuples.
Examples
julia> NNlib.scatter!(+, ones(3), [10,100], [1,3])
3-element Vector{Float64}:
11.0
1.0
101.0
julia> NNlib.scatter!(*, fill(0.5, 2, 4), [1 10; 100 1000], [3,2])
2×4 Matrix{Float64}:
0.5 5.0 0.5 0.5
0.5 500.0 50.0 0.5Sampling
NNlib.grid_sample — Function
grid_sample(input::AbstractArray{T, 4}, grid::AbstractArray{T, 4}; padding_mode = :zeros)
grid_sample(input::AbstractArray{T, 5}, grid::AbstractArray{T, 4}; padding_mode = :zeros)
Given `input`, compute output by sampling `input` values at pixel
locations from `grid`. Uses bilinear interpolation to calculate output values.
This implementation assumes the extrema (`-1` and `1`) are considered
as referring to the center points of the input’s corner pixels
(i.e. align corners is `true`).
# Arguments
- `input`: Input array in `(W_in, H_in, [D_in,] C, N)` shape.
- `grid`: Input grid in `(2, W_out, H_out, [D_out,] N)` shape.
Where for each `(W_out, H_out, [D_out,] N)` grid contains `(x, y [,z])`
coordinates that specify sampling locations normalized by the `input` shape.
Therefore, `x`, `y` and [`z`] should have values in `[-1, 1]` range.
For example, `(x = -1, y = -1, [z = -1])` is the left-top[-front] pixel of `input`,
and `(x = 1, y = 1, [z = 1])` is the right-bottom-back pixel of `input`.
Out-of-bound values are handled according to the `padding_mode`.
- `padding_mode`: Out-of-bound padding.
`:zeros` to use `0` for out-of-bound grid locations.
`:border` to use border values for out-of-bound grid locations.
Default is `:zeros`.
# Returns
`(W_out, H_out, [D_out,] C, N)` sampled grid from `input`.
# Examples
In the example below, grid contains two out-of-bound sampling locations,
which are handled differently, depending on the `padding_mode`.
```jldoctest
julia> x = reshape(collect(1.0:4.0), (2, 2, 1, 1))
2×2×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
1.0 3.0
2.0 4.0
julia> grid = Array{Float64}(undef, 2, 3, 2, 1);
julia> grid[:, 1, 1, 1] .= (-3, -1);
julia> grid[:, 2, 1, 1] .= (0, -1);
julia> grid[:, 3, 1, 1] .= (1, -1);
julia> grid[:, 1, 2, 1] .= (-1, 1);
julia> grid[:, 2, 2, 1] .= (0, 1);
julia> grid[:, 3, 2, 1] .= (3, 1);
julia> grid_sample(x, grid; padding_mode=:zeros)
3×2×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
0.0 3.0
1.5 3.5
2.0 0.0
julia> grid_sample(x, grid; padding_mode=:border)
3×2×1×1 Array{Float64, 4}:
[:, :, 1, 1] =
1.0 3.0
1.5 3.5
2.0 4.0
```NNlib.∇grid_sample — Function
∇grid_sample(Δ::AbstractArray{T, 4}, input::AbstractArray{T, 4}, grid::AbstractArray{T, 4}; padding_mode = :zeros) where TArguments
Δ: Input gradient in(W_out, H_out, C, N)shape (same as output of the primal computation).input: Input from primal computation in(W_in, H_in, C, N)shape.grid: Grid from primal computation in(2, W_out, H_out, N)shape.padding_mode: Out-of-bound padding.:zerosto use0for out-of-bound grid locations.:borderto use border values for out-of-bound grid locations. Should be the same as in primal computation. Default is:zeros.
Returns
dinput (same shape as input) and dgrid (same shape as grid) gradients.
Normalization
Flux's BatchNorm, InstanceNorm, GroupNorm and LayerNorm layers can use NNlib.batchnorm, NNlib.instancenorm, NNlib.groupnorm and NNlib.layernorm as their functional backend.
NNlib.normalise — Function
normalise(x; dims=ndims(x), eps=1f-5)Normalise x to zero mean and unit standard deviation across the dimension(s) given by dims. Per default, dims is the last dimension. eps is a small term added to the variance for numerical stability.
It is layernorm without scale and bias. For half-precision (Float16/BFloat16) inputs the statistics are computed in Float32.
Examples
julia> using Statistics
julia> x = [90, 100, 110, 130, 70];
julia> y = NNlib.normalise(x);
julia> isapprox(std(y; corrected=false), 1, atol=1e-5)
trueNNlib.batchnorm — Function
batchnorm(g, b, x, running_mean=nothing, running_var=nothing, momentum=0.1f0;
eps=1f-5, training=true, track_stats=true)Functional batch normalization. g and b are the per-channel scale and bias (either may be nothing for no affine transform); x are the feature maps. For an input with N dimensions the N-1th is the channel dimension (the usual convention for WHCN images); statistics are computed over every D_1×…×D_{N-2}×1×D_N slice, so per channel.
If running_mean/running_var are supplied:
- with
training=true(default) they are updated in place (whentrack_stats=true) with an exponential moving average controlled bymomentum, while statistics of the current batch are used to normalise; - with
training=falsethey are used to normalise (inference / test mode).
eps is added to the variance for numerical stability. On the GPU, 3D/4D/5D CuArrays are dispatched to the cuDNN implementation.
For half-precision (Float16/BFloat16) feature maps the statistics are computed in Float32 and the result is cast back; g, b, running_mean and running_var must then be Float32.
See also instancenorm, groupnorm, layernorm.
NNlib.instancenorm — Function
instancenorm(g, b, x, running_mean=nothing, running_var=nothing, momentum=0.1f0;
eps=1f-5, training=true, track_stats=false)Functional instance normalization. Arguments match batchnorm, but for an input with N > 2 dimensions statistics are computed over every D_1×…×D_{N-2}×1×1 slice, so per channel and per sample in the batch. When tracked, the running statistics (length size(x, N-1)) accumulate the per-channel average across the batch.
Normalising a dense CPU or GPU array with its own statistics, with g and b of length size(x, N-1), runs fused kernels; other cases use the generic code.
NNlib.groupnorm — Function
groupnorm(g, b, x, G::Integer; eps=1f-5)Functional group normalization. g and b are the per-channel scale and bias (either may be nothing); x are the feature maps. For an input with N > 2 dimensions the N-1th is the channel dimension; its C = size(x, N-1) channels are split into G groups (G must divide C) and statistics are computed over each group together with the spatial dimensions, per sample. eps is added to the variance.
Normalising a dense CPU or GPU array, with g and b of length C, runs fused kernels; other cases use the generic code.
See also batchnorm, instancenorm, layernorm.
NNlib.layernorm — Function
layernorm(g, b, x; dims=1, eps=1f-5)Functional layer normalization. g and b are the scale and bias (either may be nothing); x is normalised over the dimensions dims (the leading dimension by default).
g and b, when given, must broadcast against the normalised region, e.g. have size size(x)[dims] with singleton trailing dimensions.
Normalising a dense CPU or GPU array over leading dims (1:k), with g and b of the shape of those dimensions, runs fused kernels; other cases use the generic code.
See also normalise, batchnorm, instancenorm, groupnorm.
NNlib.rmsnorm — Function
rmsnorm(g, x; dims=1, eps=1f-5)Functional root mean square normalization. x is divided by its root mean square over the dimensions dims (the leading dimension by default) and multiplied by the scale g (nothing for no scaling):
y = g .* x ./ sqrt.(mean(abs2, x; dims) .+ eps)Unlike layernorm, the mean is not subtracted and there is no bias.
g, when given, must broadcast against the normalised region, e.g. have size size(x)[dims] with singleton trailing dimensions. eps is added to the mean square for numerical stability.
For half-precision (Float16/BFloat16) inputs the mean square is computed in Float32 and the result is cast back; g must then be Float32.
See also layernorm, normalise.
Examples
julia> x = Float32[3 0; 4 2];
julia> NNlib.rmsnorm(nothing, x) ≈ x ./ sqrt.(Float32[12.5 2])
true
julia> NNlib.rmsnorm(Float32[2, 1], x) ≈ Float32[2, 1] .* NNlib.rmsnorm(nothing, x)
trueNNlib.∇batchnorm — Function
∇batchnorm(g, b, x, dy, running_mean=nothing, running_var=nothing, momentum=0.1f0;
eps=1f-5, training=true)Gradient of batchnorm: given the upstream gradient dy, return (dg, db, dx). On the GPU, 3D/4D/5D CuArrays are dispatched to the cuDNN implementation in NNlibCUDACUDNNExt; this generic method is the fallback used everywhere else and is itself differentiable (second-order).
NNlib.∇instancenorm — Function
∇instancenorm(g, b, x, dy, running_mean=nothing, running_var=nothing, momentum=0.1f0;
eps=1f-5, training=true)Gradient of instancenorm, returning (dg, db, dx).
NNlib.∇groupnorm — Function
∇groupnorm(g, b, x, dy, G::Integer; eps=1f-5)Gradient of groupnorm, returning (dg, db, dx).
NNlib.∇layernorm — Function
∇layernorm(g, b, x, dy; dims=1, eps=1f-5)Gradient of layernorm, returning (dg, db, dx).
NNlib.∇rmsnorm — Function
∇rmsnorm(g, x, dy; dims=1, eps=1f-5)Gradient of rmsnorm, returning (dg, dx).
Losses
NNlib.ctc_loss — Function
ctc_loss(ŷ, y)Computes the connectionist temporal classification loss between ŷ and y. ŷ must be a classes-by-time matrices, i.e., each row represents a class and each column represents a time step. Additionally, the logsoftmax function will be applied to ŷ, so ŷ must be the raw activation values from the neural network and not, for example, the activations after being passed through a softmax activation function. y must be a 1D array of the labels associated with ŷ. The blank label is assumed to be the last label category in ŷ, so it is equivalent to size(ŷ, 1). Used for sequence-to-sequence classification problems such as speech recognition and handwriting recognition where the exact time-alignment of the output (e.g., letters) is not needed to solve the problem. See Graves et al. (2006) or Graves (2012) for mathematical details.
Miscellaneous
NNlib.logsumexp — Function
logsumexp(x; dims = :)Computes log.(sum(exp.(x); dims)) in a numerically stable way. Without dims keyword this returns a scalar. For Float16 and BFloat16 inputs, it is computed and returned in Float32.
See also logsoftmax.
NNlib.glu — Function
glu(x, dim = 1)The gated linear unit from the "Language Modeling with Gated Convolutional Networks" paper.
Calculates a .* sigmoid(b), where x is split in half along given dimension dim to form a and b.
See also swiglu.
NNlib.swiglu — Function
swiglu(x, dim = 1)The swish-gated linear unit from the "GLU Variants Improve Transformer" paper, the feed-forward activation of Llama, PaLM and most language models since.
Calculates swish.(gate) .* up, where x is split in half along given dimension dim to form gate (the first half) and up (the second half). This is the layout of the fused gate and up projections of Llama-style models, and the reverse of glu, where the gate is the second half.
See also swish.
Examples
ffn = Chain(Dense(d => 2h), NNlib.swiglu, Dense(h => d))NNlib.within_gradient — Function
within_gradient(x) --> BoolReturns false except when used inside a gradient call, when it returns true. Useful for Flux regularisation layers which behave differently during training and inference.
This should work with any ChainRules-based differentiation package, in which case x is ignored. But Tracker.jl overloads within_gradient(x::TrackedArray), thus for widest use you should pass it an array whose gradient is of interest. There is also an overload for ForwardDiff.jl's Dual types (and arrays of them). Enzyme.jl and Mooncake.jl are supported through package extensions, and ignore x too.
Examples
julia> using ForwardDiff, Zygote, NNlib
julia> f_good(x) = if NNlib.within_gradient(x)
@show 10x
else
x
end;
julia> Zygote.withgradient(f_good, 1.0)
10x = 10.0
(val = 10.0, grad = (10.0,))
julia> ForwardDiff.derivative(f_good, 1.0)
10x = Dual{ForwardDiff.Tag{typeof(f_good), Float64}}(10.0,10.0)
10.0
julia> f_bad(x, y) = if any(NNlib.within_gradient, (x, y))
@show x * y
else
x / y
end;
julia> Zygote.withgradient(f_bad, 2.0, 3.0)
(val = 0.6666666666666666, grad = (0.3333333333333333, -0.2222222222222222))
julia> ForwardDiff.derivative(x -> f_bad(x, 3.0), 2.0)
x * y = Dual{ForwardDiff.Tag{var"#9#10", Float64}}(6.0,3.0)
3.0What goes wrong in f_bad is that Zygote knows any to be non-differentiable, and thus completely ignores its contents. This is not a perfect mechanism, and the only style recommended is precisely that of f_good above.
NNlib.bias_act! — Function
bias_act!(σ, x, b)This is equivalent to x .= σ.(x .+ b), also replacing sigmoid & tanh with sigmoid_fast & tanh_fast. It will only overwrite x when x isa StridedArray{<:AbstractFloat}.
When used within a gradient, it will overwrite only when σ has a method of derivatives_given_output which does not need the input at all. Such methods are defined by e.g. @scalar_rule relu(x) Ω > 0 where the derivative contains only Ω (the output) not x.