xlstm_jax.models.xlstm_clean.blocks.xlstm_block

Classes

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

Module Contents

class xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlockConfig

mlstm: xlstm_jax.models.xlstm_clean.blocks.mlstm.layer.mLSTMLayerConfig | None = None

slstm: None = None

feedforward: xlstm_jax.models.xlstm_clean.components.feedforward.FeedForwardConfig | None = None

dtype: str = 'bfloat16'

_num_blocks: int | None = None

_block_idx: int | None = None

property _dtype: jax.numpy.dtype

Returns the real dtype instead of the str from configs.

Returns:

The jnp dtype corresponding to the string value.

Return type:

jax.numpy.dtype

class xlstm_jax.models.xlstm_clean.blocks.xlstm_block.xLSTMBlock

Bases: flax.linen.Module

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

It contains the pre-LayerNorms and the skip connections.

config: xLSTMBlockConfig