aihwkit.nn.modules.rnn.layers module

Analog RNN layers

class aihwkit.nn.modules.rnn.layers.AnalogBidirRNNLayer(cell, *cell_args)

Bases: Module

Bi-directional analog RNN layer.

Parameters:

• cell (Type) – RNNCell type (AnalogLSTMCell/AnalogGRUCell/AnalogVanillaRNNCell)

• cell_args (Any) – arguments to RNNCell (e.g. input_size, hidden_size, rpu_configs)

forward(input_, states)

Forward pass.

Parameters:

• input – input tensor

• states (List[Tuple[Tensor, Tensor] | Tensor]) – LSTM state tensor

• input_ (Tensor) –

Returns:

cat outputs and states

Return type:

Tuple[Tensor, List[Tuple[Tensor, Tensor] | Tensor]]

get_zero_state(batch_size)

Returns a zeroed state.

Parameters:

batch_size (int) – batch size of the input

Returns:

Zeroed state tensor

Return type:

Tensor

class aihwkit.nn.modules.rnn.layers.AnalogRNNLayer(cell, *cell_args)

Bases: Module

Analog RNN Layer.

Parameters:

• cell (Type) – RNNCell type (AnalogLSTMCell/AnalogGRUCell/AnalogVanillaRNNCell/ AnalogLSTMCellSingleRPU)

• cell_args (Any) – arguments to RNNCell (e.g. input_size, hidden_size, rpu_configs)

forward(input_, state)

Forward pass.

Parameters:

• input – input tensor

• state (Tuple[Tensor, Tensor] | Tensor) – LSTM state tensor

• input_ (Tensor) –

Returns:

stacked outputs and state

Return type:

Tuple[Tensor, Tuple[Tensor, Tensor]]

get_zero_state(batch_size)

Returns a zeroed state.

Parameters:

batch_size (int) – batch size of the input

Returns:

Zeroed state tensor

Return type:

Tensor

class aihwkit.nn.modules.rnn.layers.AnalogReverseRNNLayer(cell, *cell_args)

Bases: Module

Analog RNN layer for direction.

Parameters:

• cell (Type) – RNNCell type (AnalogLSTMCell/AnalogGRUCell/AnalogVanillaRNNCell)

• cell_args (Any) – arguments to RNNCell (e.g. input_size, hidden_size, rpu_configs)

forward(input_, state)

Forward pass.

Parameters:

• input – input tensor

• state (Tuple[Tensor, Tensor] | Tensor) – LSTM state tensor

• input_ (Tensor) –

Returns:

stacked reverse outputs and state

Return type:

Tuple[Tensor, Tuple[Tensor, Tensor] | Tensor]

get_zero_state(batch_size)

Returns a zeroed state.

Parameters:

batch_size (int) – batch size of the input

Returns:

Zeroed state tensor

Return type:

Tensor

static reverse(lst)

Reverses the list of input tensors.

Parameters:

lst (List[Tensor]) –

Return type:

List[Tensor]