xlstm_jax.models.xlstm_pytorch.blocks.slstm.block

Classes

sLSTMBlockConfig
sLSTMBlock	An xLSTM block can be either an sLSTM Block or an mLSTM Block.

Module Contents

class xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlockConfig

slstm: xlstm_jax.models.xlstm_pytorch.blocks.slstm.layer.sLSTMLayerConfig

feedforward: xlstm_jax.models.xlstm_pytorch.components.feedforward.FeedForwardConfig | None

_num_blocks: int | None = None

_block_idx: int | None = None

class xlstm_jax.models.xlstm_pytorch.blocks.slstm.block.sLSTMBlock(config)

Bases: xlstm_jax.models.xlstm_pytorch.blocks.xlstm_block.xLSTMBlock

An xLSTM block can be either an sLSTM Block or an mLSTM Block.

It contains the pre-LayerNorms and the skip connections.

Parameters:

config (sLSTMBlockConfig)

config_class

config

xlstm_norm

forward(x, **kwargs)

Parameters:

x (torch.Tensor)

Return type:

torch.Tensor

step(x, **kwargs)

Parameters:

x (torch.Tensor)

Return type:

tuple[torch.Tensor, dict[str, tuple[torch.Tensor, Ellipsis]]]

reset_parameters()

Return type:

None