aihwkit.nn.modules.lstm module

Analog LSTM layers.

class aihwkit.nn.modules.lstm.AnalogLSTM(input_size, hidden_size, num_layers=1, dropout=0.0, bias=True, rpu_config=None, realistic_read_write=False, weight_scaling_omega=0.0, xavier=False)

Bases: aihwkit.nn.modules.container.AnalogSequential

Modular LSTM that uses analog tiles.

Parameters

• input_size – in_features to W_{ih} matrix of first layer

• hidden_size – in_features and out_features for W_{hh} matrices

• num_layers – number of serially connected LSTM layers

• dropout – dropout applied to output of all LSTM layers except last

• bias – whether to use a bias row on the analog tile or not

• rpu_config – resistive processing unit configuration.

• realistic_read_write – whether to enable realistic read/write for setting initial weights and read out of weights

• weight_scaling_omega – the weight value where the max weight will be scaled to. If zero, no weight scaling will be performed

• xavier – whether standard PyTorch LSTM weight initialization (default) or Xavier initialization

forward(x, states=None)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters

• x (torch.Tensor) –

• states (Optional[List[Tuple[torch.Tensor, torch.Tensor]]]) –

Return type

Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]

init_layers(weight_init_fn, bias_init_fn=None)

Init the analog layers with custom functions.

Parameters

• weight_init_fn (Callable) – in-place tensor function applied to weight of AnalogLinear layers

• bias_init_fn (Optional[Callable]) – in-place tensor function applied to bias of AnalogLinear layers

Return type

None

Note

If no bias init function is provided the weight init function is taken for the bias as well.

reset_parameters(xavier=False)

Weight and bias initialization.

Parameters

xavier (bool) – whether standard PyTorch LSTM weight initialization (default) or Xavier initialization

Return type

None

class aihwkit.nn.modules.lstm.AnalogLSTMCell(input_size, hidden_size, bias, rpu_config, realistic_read_write=False, weight_scaling_omega=0.0)

Bases: aihwkit.nn.modules.container.AnalogSequential

Analog LSTM Cell.

Parameters

• input_size – in_features size for W_ih matrix

• hidden_size – in_features and out_features size for W_hh matrix

• bias – whether to use a bias row on the analog tile or not

• rpu_config – configuration for an analog resistive processing unit

• realistic_read_write – whether to enable realistic read/write for setting initial weights and read out of weights

• weight_scaling_omega – the weight value where the max weight will be scaled to. If zero, no weight scaling will be performed

forward(input_, state)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters

• input_ (torch.Tensor) –

• state (Tuple[torch.Tensor, torch.Tensor]) –

Return type

Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

class aihwkit.nn.modules.lstm.AnalogLSTMLayer(cell, *cell_args)

Bases: aihwkit.nn.modules.container.AnalogSequential

Analog LSTM Layer.

Parameters

• cell – LSTMCell type (e.g. AnalogLSTMCell)

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

forward(input_, state)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters

• input_ (torch.Tensor) –

• state (Tuple[torch.Tensor, torch.Tensor]) –

Return type

Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]

class aihwkit.nn.modules.lstm.LSTMState(hx, cx)

Bases: tuple

property cx

Alias for field number 1

property hx

Alias for field number 0

class aihwkit.nn.modules.lstm.ModularAnalogLSTMWithDropout(num_layers, layer, dropout, first_layer_args, other_layer_args)

Bases: aihwkit.nn.modules.container.AnalogSequential

Modular Analog LSTM with dropout.

Parameters

• num_layers – number of serially connected LSTM layers

• layer – LSTM layer type (e.g. AnalogLSTMLayer)

• dropout – dropout applied to output of all LSTM layers except last

• first_layer_args – LSTMCell type, input_size, hidden_size, rpu_config, etc.

• other_layer_args – LSTMCell type, hidden_size, hidden_size, rpu_config, etc.

forward(input_, states)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters

• input_ (torch.Tensor) –

• states (List[Tuple[torch.Tensor, torch.Tensor]]) –

Return type

Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]

aihwkit.nn.modules.lstm.init_stacked_analog_lstm(num_layers, layer, first_layer_args, other_layer_args)

Construct a list of LSTMLayers over which to iterate.

Parameters

• num_layers (int) – number of serially connected LSTM layers

• layer (Type) – LSTM layer type (e.g. AnalogLSTMLayer)

• first_layer_args (Any) – LSTMCell type, input_size, hidden_size, rpu_config, etc.

• other_layer_args (Any) – LSTMCell type, hidden_size, hidden_size, rpu_config, etc.

Returns

torch.nn.ModuleList, which is similar to a regular Python list, but where torch.nn.Module methods can be applied

Return type

torch.nn.modules.container.ModuleList