xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend

Submodules

• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.attention
• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config
• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config_utils
• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.fwbw
• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.layer_factory
• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent
• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.recurrent_triton
• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.simple
• xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.triton_kernels

Attributes

_mlstm_backend_registry
BackendType
create_mlstm_backend

Classes

mLSTMBackendNameAndKwargs

Sub-model configuration.

Package Contents

xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend._mlstm_backend_registry

xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.BackendType

class xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.mLSTMBackendNameAndKwargs

Bases: config_utils.NameAndKwargs

Sub-model configuration.

This class is currently a quick fix to allow for post-init style model configs, like the xlstm-clean we ported from the original xlstm codebase. Once the config system is more mature, we should remove this and all becomes a subclass of ModelConfig.

_registry: dict[str, type]

name: str = ''

kwargs: dict[str, Any] | None

get_config_class_for_kwargs()

Return type:

type

get_class_for_name()

Return type:

type

create_name_instance()

Return type:

Any

to_dict()

Converts the config to a dictionary.

Helpful for saving to disk or logging.

Return type:

dict

xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.create_mlstm_backend