xlstm_jax.models.xlstm_parallel.utils#

Classes#

UpProjConfigMixin

Sub-model configuration.

Module Contents#

class xlstm_jax.models.xlstm_parallel.utils.UpProjConfigMixin#

Bases: xlstm_jax.models.configs.SubModelConfig

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.

proj_factor: float | None = None#

round_proj_up_dim_up: bool = True#

round_proj_up_to_multiple_of: int = 64#

_proj_up_dim: int | None = None#

_set_proj_up_dim(embedding_dim)#

Parameters:: embedding_dim (int)
Return type:: None

to_dict()#

Converts the config to a dictionary.

Helpful for saving to disk or logging.

Return type:: dict