Implementation:LaurentMazare Tch rs RNN

Knowledge Sources	LaurentMazare_Tch_rs
Domains	Neural Networks, Recurrent Networks, Deep Learning
Last Updated	2026-02-08 00:00 GMT

Overview

The rnn module implements LSTM and GRU recurrent neural network layers, unified under a common RNN trait that provides zero_state, step, seq, and seq_init methods for sequence processing.

Description

This module provides a complete recurrent neural network abstraction. The RNN trait defines an associated State type and four methods: zero_state creates an initial hidden state, step processes a single time step, seq runs over a full sequence from a zero state, and seq_init runs over a full sequence from a given initial state. The seq method has a default implementation that calls zero_state then seq_init.

LSTM (Long Short-Term Memory) uses a gated architecture with four gates (gate_dim = 4 * hidden_dim). Its state, LSTMState, wraps a tuple of two tensors: the hidden state h and cell state c, each accessible via dedicated methods. The lstm constructor creates the layer by allocating flat weight tensors and optionally flattening them for cuDNN optimization when on CUDA.

GRU (Gated Recurrent Unit) uses three gates (gate_dim = 3 * hidden_dim). Its state, GRUState, wraps a single hidden tensor accessible via value(). The gru constructor follows the same pattern as LSTM.

Both layers share RNNConfig, which configures: has_biases, num_layers, dropout, train, bidirectional, batch_first, and weight/bias initialization strategies (w_ih_init, w_hh_init, b_ih_init, b_hh_init). The internal rnn_weights function allocates all weight and bias tensors for all layers and directions.

Usage

Use LSTM for tasks requiring long-range dependencies with explicit memory control. Use GRU for a simpler recurrent model with fewer parameters. Both support multi-layer stacking, bidirectional processing, and dropout between layers.

Code Reference

Source Location

• Repository: LaurentMazare_Tch_rs
• File: src/nn/rnn.rs

Signature

// RNN trait
pub trait RNN {
    type State;
    fn zero_state(&self, batch_dim: i64) -> Self::State;
    fn step(&self, input: &Tensor, state: &Self::State) -> Self::State;
    fn seq(&self, input: &Tensor) -> (Tensor, Self::State);
    fn seq_init(&self, input: &Tensor, state: &Self::State) -> (Tensor, Self::State);
}

// LSTM
pub struct LSTMState(pub (Tensor, Tensor));
impl LSTMState {
    pub fn h(&self) -> Tensor;
    pub fn c(&self) -> Tensor;
}

pub struct LSTM {
    flat_weights: Vec<Tensor>,
    hidden_dim: i64,
    config: RNNConfig,
    device: Device,
}
pub fn lstm(vs: T, in_dim: i64, hidden_dim: i64, c: RNNConfig) -> LSTM;

// GRU
pub struct GRUState(pub Tensor);
impl GRUState {
    pub fn value(&self) -> Tensor;
}

pub struct GRU {
    flat_weights: Vec<Tensor>,
    hidden_dim: i64,
    config: RNNConfig,
    device: Device,
}
pub fn gru(vs: T, in_dim: i64, hidden_dim: i64, c: RNNConfig) -> GRU;

// Shared config
pub struct RNNConfig {
    pub has_biases: bool,
    pub num_layers: i64,
    pub dropout: f64,
    pub train: bool,
    pub bidirectional: bool,
    pub batch_first: bool,
    pub w_ih_init: super::Init,
    pub w_hh_init: super::Init,
    pub b_ih_init: Option<super::Init>,
    pub b_hh_init: Option<super::Init>,
}

Import

use tch::nn::{lstm, gru, RNNConfig, RNN, LSTM, GRU, LSTMState, GRUState};

I/O Contract

Constructor Parameter	Type	Description
vs	impl Borrow<Path>	Variable store path for parameter allocation
in_dim	i64	Input feature dimension
hidden_dim	i64	Hidden state dimension
c	RNNConfig	Shared configuration for LSTM/GRU

RNNConfig Field	Default	Description
has_biases	true	Include bias terms
num_layers	1	Number of stacked recurrent layers
dropout	0.0	Dropout between layers (not applied to last layer)
train	true	Training mode flag
bidirectional	false	Process input in both directions
batch_first	true	Input shape is [batch, seq_len, features]
w_ih_init	DEFAULT_KAIMING_UNIFORM	Input-hidden weight init
w_hh_init	DEFAULT_KAIMING_UNIFORM	Hidden-hidden weight init
b_ih_init	Some(Const(0.))	Input-hidden bias init
b_hh_init	Some(Const(0.))	Hidden-hidden bias init

RNN Method	Input	Output
zero_state(batch_dim)	Batch size	Initial state (zeros)
step(input, state)	[batch, features] tensor + state	Updated state
seq(input)	[batch, seq_len, features] tensor	(output tensor, final state)
seq_init(input, state)	[batch, seq_len, features] tensor + initial state	(output tensor, final state)

Usage Examples

use tch::{nn, nn::RNN, Device, Kind, Tensor};

let vs = nn::VarStore::new(Device::Cpu);
let root = vs.root();

// Create a 2-layer bidirectional LSTM
let lstm_layer = nn::lstm(
    &root / "lstm",
    128,  // input features
    256,  // hidden size
    nn::RNNConfig { num_layers: 2, bidirectional: true, ..Default::default() },
);

// Process a sequence
let input = Tensor::randn([4, 20, 128], (Kind::Float, Device::Cpu));
let (output, state) = lstm_layer.seq(&input);
// output shape: [4, 20, 512] (256 * 2 for bidirectional)
let hidden = state.h();
let cell = state.c();

// Create a GRU and process step-by-step
let gru_layer = nn::gru(&root / "gru", 64, 128, Default::default());
let mut state = gru_layer.zero_state(4);
for t in 0..10 {
    let step_input = Tensor::randn([4, 64], (Kind::Float, Device::Cpu));
    state = gru_layer.step(&step_input, &state);
}
let final_hidden = state.value();

Related Pages

• Principle:LaurentMazare_Tch_rs_Recurrent_Neural_Network