nn.SO3GatedNonlinearity

class nn.SO3GatedNonlinearity(*args: Any, **kwargs: Any)

SO3-equivariant gated nonlinearity.

With combined indexing s=(l,m), this can be written as:

\[y_{i,s,f} = x_{j,s,f} * \sigma(x_{j,0,\cdot})\]