xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block

Classes

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

Module Contents

class xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlockConfig

mlstm: xlstm_jax.models.xlstm_pytorch.blocks.mlstm.layer.mLSTMLayerConfig

_num_blocks: int | None = None

_block_idx: int | None = None

class xlstm_jax.models.xlstm_pytorch.blocks.mlstm.block.mLSTMBlock(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 (mLSTMBlockConfig)

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