Activations

Non-linear activations, available both as stateful layer modules and as pure functions for flexible composition.

Element-wise Layers

Non-linear activation layers that operate element-wise on input tensors. Includes rectified linear units (ReLU and variants), sigmoid functions, hyperbolic tangent, softmax for probability distributions, and specialized activations for specific architectures (SELU, GELU, SiLU, Mish). These introduce non-linearity enabling networks to learn complex patterns.

Threshold	Thresholds each element of the input Tensor.
ReLU	Applies the rectified linear unit function element-wise.
RReLU	Applies the randomized leaky rectified liner unit function, element-wise.
Hardtanh	Applies the HardTanh function element-wise.
ReLU6	Applies the element-wise function.
Sigmoid	Applies the element-wise function.
Hardsigmoid	Applies the Hardsigmoid function element-wise.
Tanh	Applies the Hyperbolic Tangent (Tanh) function element-wise.
SiLU	Applies the Sigmoid Linear Unit (SiLU) function, element-wise.
Mish	Applies the Mish function, element-wise.
Hardswish	Applies the Hardswish function, element-wise.
ELU	Applies the Exponential Linear Unit (ELU) function, element-wise.
CELU	Applies the element-wise function.
SELU	Applied element-wise.
GLU	Applies the gated linear unit function.
GELU	Applies the Gaussian Error Linear Units function.
Hardshrink	Applies the Hard Shrinkage (Hardshrink) function element-wise.
LeakyReLU	Applies the element-wise function.
LogSigmoid	Applies the element-wise function.
Softplus	Applies the Softplus function element-wise.
Softshrink	Applies the soft shrinkage function elementwise.
PReLU	Applies the element-wise function.
Softsign	Applies the element-wise function.
Tanhshrink	Applies the element-wise function.
Softmin	Applies the Softmin function to an n-dimensional input Tensor.
Softmax	Applies the Softmax function to an n-dimensional input Tensor.
Softmax2d	Applies SoftMax over features to each spatial location.
LogSoftmax	Applies the \(\log(\text{Softmax}(x))\) function to an n-dimensional input Tensor.
Identity	A placeholder identity operator that is argument-insensitive.
SpikeBitwise	Bitwise addition for the spiking inputs.

Functional Activations

Functional (non-module) activation functions for flexible composition. These are pure functions that can be used directly in update() methods or combined with JAX transformations. Provides the same activations as the layer-based equivalents but without state or module overhead.

tanh	Hyperbolic tangent activation function.
relu	Rectified Linear Unit activation function.
squareplus	Squareplus activation function.
softplus	Softplus activation function.
soft_sign	Soft-sign activation function.
sigmoid	Sigmoid activation function.
silu	SiLU (Sigmoid Linear Unit) activation function.
swish	SiLU (Sigmoid Linear Unit) activation function.
log_sigmoid	Log-sigmoid activation function.
elu	Exponential Linear Unit activation function.
leaky_relu	Leaky Rectified Linear Unit activation function.
hard_tanh	Hard hyperbolic tangent activation function.
celu	Continuously-differentiable Exponential Linear Unit activation.
selu	Scaled Exponential Linear Unit activation.
gelu	Gaussian Error Linear Unit activation function.
glu	Gated Linear Unit activation function.
logsumexp
log_softmax	Log-Softmax function.
softmax	Softmax activation function.
standardize	Standardize (normalize) an array.
one_hot	One-hot encode the given indices.
relu6	Rectified Linear Unit 6 activation function.
hard_sigmoid	Hard Sigmoid activation function.
hard_silu	Hard SiLU (Swish) activation function.
hard_swish	Hard SiLU (Swish) activation function.
hard_shrink	Hard shrinkage activation function.
rrelu	Randomized Leaky Rectified Linear Unit activation function.
mish	Mish activation function.
soft_shrink	Soft shrinkage activation function.
prelu	Parametric Rectified Linear Unit activation function.
tanh_shrink	Tanh shrink activation function.
softmin	Softmin activation function.
sparse_plus	Sparse plus activation function.
sparse_sigmoid	Sparse sigmoid activation function.