Layers

Embedding(num_embeddings, dims)	Implements a simple lookup table that maps each input integer to a high-dimensional vector.
ReLU()	Applies the Rectified Linear Unit.
PReLU([num_parameters, init])
GELU([approx])	Applies the Gaussian Error Linear Units.
SiLU()	Applies the Sigmoid Linear Unit.
Step([threshold])	Applies the Step Activation Function.
SELU()	Applies the Scaled Exponential Linear Unit.
Mish()	Applies the Mish function, element-wise.
Linear(input_dims, output_dims[, bias])	Applies an affine transformation to the input.
Conv1d(in_channels, out_channels, kernel_size)	Applies a 1-dimensional convolution over the multi-channel input sequence.
Conv2d(in_channels, out_channels, kernel_size)	Applies a 2-dimensional convolution over the multi-channel input image.
LayerNorm(dims[, eps, affine])	Applies layer normalization [1] on the inputs.
RMSNorm(dims[, eps])	Applies Root Mean Square normalization [1] to the inputs.
GroupNorm(num_groups, dims[, eps, affine, ...])	Applies Group Normalization [1] to the inputs.
RoPE(dims[, traditional, base])	Implements the rotary positional encoding [1].
MultiHeadAttention(dims, num_heads[, ...])	Implements the scaled dot product attention with multiple heads.
Sequential(*modules)	A layer that calls the passed callables in order.