e3x.nn.activations

Equivariant implementation of common activation functions.

When working with equivariant features, any nonlinear operation may only be applied to the scalar channel, else the features lose their equivariance. Consequently, standard activation functions may not be applied naively to equivariant features. However, it is possible to apply gated linear activations

\[\mathrm{gated\_linear}(\mathbf{x}) = \mathrm{gate}\left(\mathbf{x}^{(0_+)}\right) \circ \mathbf{x}\]

where the nonlinear \(\mathrm{gate}\) function is applied element-wise on the scalar channel \(\mathbf{x}^{(0_+)}\) (the element-wise product ‘\(\circ\)’ implies broadcasting over dimensions). Many common activation functions can be written as (equivariant) gated linear activations, such that for scalar inputs, they are equivalent to their non-equivariant counterpart. See also here for more details on the notation and here for features in e3x in general.

Functions

bent_identity(x)

Bent identity activation function.

celu(x[, alpha])

Continuously-differentiable exponential linear unit activation.

elu(x[, alpha])

Exponential linear unit activation function.

gelu(x[, approximate])

Gaussian error linear unit activation function.

hard_silu(x)

Hard SiLU/Swish activation function.

hard_swish(x)

Hard SiLU/Swish activation function.

hard_tanh(x)

Hard \(\mathrm{tanh}\) activation function.

leaky_relu(x[, negative_slope])

Leaky rectified linear unit activation function.

mish(x[, fast])

Mish activation function.

relu(x)

Rectified linear unit activation function.

relu6(x)

Rectified Linear Unit 6 activation function.

selu(x)

Scaled exponential linear unit activation.

serf(x)

Serf activation function.

shifted_softplus(x)

Shifted-softplus activation function.

silu(x)

SiLU/Swish activation function.

soft_sign(x)

Soft-sign activation function.

swish(x)

SiLU/Swish activation function.