Skip to main content

Module modules

Module modules 

Expand description

Neural network modules implementations.

Modules§

attention
Attention module
basic
Basic RNN.
cache
Cache module
conv
Convolution module
gru
Gated Recurrent Unit module.
interpolate
Interpolate module
lstm
Long Short-Term Memory module. LSTM module.
norm
Normalization Layers
pool
Pooling module
transformer
Transformer module

Structs§

BatchNorm
Applies Batch Normalization over a tensor.
BatchNormConfig
BatchNorm Configuration.
BiGru
The BiGru module. This implementation is for Bidirectional GRU.
BiGruConfig
Configuration to create a BiGru module using the init function.
BiLstm
The BiLstm module. This implementation is for Bidirectional LSTM.
BiLstmConfig
Configuration to create a BiLstm module using the init function.
BiRnn
The BiRnn module. This implementation is for Bidirectional RNN. Should be created with BiRnnConfig.
BiRnnConfig
Configuration to create a BiRnn module using the init function.
CosineSimilarity
Computes the cosine similarity between two tensors along a dimension, following torch.nn.CosineSimilarity.
CosineSimilarityConfig
Configuration to create a CosineSimilarity layer using the init function.
Dropout
Set at random some elements of the input tensor to zero during training.
DropoutConfig
Configuration to create a Dropout layer using the init function.
Embedding
Lookup table to store a fix number of vectors.
EmbeddingConfig
Configuration to create an Embedding layer using the init function.
Fold4d
Four-dimensional folding, the inverse of Unfold4d.
Fold4dConfig
Configuration to create a fold 4d layer using the init function.
GateController
A GateController represents a gate in an LSTM cell. An LSTM cell generally contains three gates: an input gate, forget gate, and output gate. Additionally, cell gate is just used to compute the cell state.
GaussianNoise
Add pseudorandom Gaussian noise to an arbitrarily shaped tensor.
GaussianNoiseConfig
Configuration to create a GaussianNoise layer using the init function.
GroupNorm
Applies Group Normalization over a mini-batch of inputs as described in the paper Group Normalization.
GroupNormConfig
Configuration to create a GroupNorm layer using the init function.
Gru
The Gru (Gated recurrent unit) module. This implementation is for a unidirectional, stateless, Gru.
GruConfig
Configuration to create a gru module using the init function.
Identity
Identity Module.
InstanceNorm
Applies Instance Normalization over a tensor as described in the paper Instance Normalization
InstanceNormConfig
Configuration to create a InstanceNorm layer using the init function.
LayerNorm
Applies Layer Normalization over an input tensor as described in the paper Layer Normalization.
LayerNormConfig
Configuration to create a LayerNorm layer using the init function.
Linear
Applies a linear transformation to the input tensor.
LinearConfig
Configuration to create a Linear layer using the init function.
LocalResponseNorm
Applies Local Response Normalization as described in ImageNet Classification with Deep Convolutional Neural Networks.
LocalResponseNormConfig
Configuration to create a LocalResponseNorm layer using the init function.
Lstm
The Lstm module. This implementation is for a unidirectional, stateless, Lstm.
LstmConfig
Configuration to create a Lstm module using the init function.
LstmState
State bundle for LSTM implementations.
PairwiseDistance
Computes the pairwise distance between the rows of two tensors, following torch.nn.PairwiseDistance.
PairwiseDistanceConfig
Configuration to create a PairwiseDistance layer using the init function.
PixelShuffle
Rearranges a tensor of shape [N, C * r^2, H, W] into [N, C, H * r, W * r], where r is the upscale factor. Inverse of PixelUnshuffle, matching torch.nn.PixelShuffle.
PixelShuffleConfig
Configuration to create a PixelShuffle layer using the init function.
PixelUnshuffle
Rearranges a tensor of shape [N, C, H * r, W * r] into [N, C * r^2, H, W], where r is the downscale factor. Inverse of PixelShuffle, matching torch.nn.PixelUnshuffle.
PixelUnshuffleConfig
Configuration to create a PixelUnshuffle layer using the init function.
PositionalEncoding
Positional encoding layer for transformer models.
PositionalEncodingConfig
Configuration to create a PositionalEncoding layer using the init function.
RmsNorm
Applies RMS Normalization over an input tensor along the last dimension.
RmsNormConfig
Configuration to create a RMS Norm layer using the init function.
Rnn
The Rnn module. This implementation is for a unidirectional, stateless, Rnn. Should be created with RnnConfig.
RnnConfig
Configuration to create a Rnn module using the init function.
RnnState
A RnnState is used to store hidden state in RNN.
RotaryEncoding
A module that applies rotary positional encoding to a tensor. Rotary Position Encoding or Embedding (RoPE), is a type of position embedding which encodes absolute positional information with rotation matrix and naturally incorporates explicit relative position dependency in self-attention formulation.
RotaryEncodingConfig
Configuration to create a RotaryEncoding layer using the init function.
Unfold4d
Four-dimensional unfolding.
Unfold4dConfig
Configuration to create an unfold 4d layer using the init function.

Enums§

LinearLayout
The layout in which the linear parameters are stored.

Traits§

OptionalInitialLstmState
Extension trait for attaching an initializer to an Option<LstmState>.

Functions§

generate_sinusoids
Returns sinusoids for positional embedding introduced in Attention is all you need.