xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config

Classes

mLSTMBackend

Module Contents

class xlstm_jax.models.xlstm_parallel.blocks.mlstm.backend.config.mLSTMBackend

Bases: flax.linen.Module

config: Any

property can_vmap_over_heads: bool

Abstractmethod:

Return type:

bool

Whether the backend can be vmaped over the heads dimension.

Should be False if the backend requires manual transposition of the input tensors to work over heads.