XLSTM

class neuralhydrology.modelzoo.x_lstm.XLSTM(cfg: Config)

Bases: BaseModel

Extended Long Short-Term Memory (xLSTM) backbone model class, which relies on https://github.com/NX-AI/xlstm

This class implements the xLSTM backbone, consisting of stacked mLSTM and sLSTM blocks. A transition layer (linear projection) ensures that the input feature dimension matches the hidden size required by the xLSTM stack.

Parameters:

cfg (Config) – The run configuration.

forward(data: dict[str, Tensor | dict[str, Tensor]]) → Dict[str, Tensor]

Perform a forward pass on the xlstm model.

Parameters:

data (dict[str, torch.Tensor | dict[str, torch.Tensor]]) – Dictionary, containing input features as key-value pairs.

Returns:

• ‘y_hat’: model predictions of shape [batch size, sequence length, number of target variables].

Return type:

Dict[str, torch.Tensor]

module_parts = ['embedding_net', 'transition_layer', 'xlstm', 'dropout', 'head']