def test_model_generate(
    batch_size: int,
    prefill_length: int,
    max_length: int,
    chunkwise_kernel_name: str = "chunkwise--native_autograd",
    sequence_kernel_name: str = "native_sequence__native",
    step_kernel_name: str = "native",
    seed: int = 0,
    vocab_size: int = 512,
    num_blocks: int = 2,
    embedding_dim: int = 256,
    num_heads: int = 4,
    device: str = "cuda",
):
    """Tests calling the generation from the model."""
    torch.manual_seed(seed)
    device = torch.device(device)
    mlstm_config = xLSTMLargeConfig(
        embedding_dim=embedding_dim,
        num_heads=num_heads,
        num_blocks=num_blocks,
        vocab_size=vocab_size,
        mode="inference",
        chunkwise_kernel=chunkwise_kernel_name,
        sequence_kernel=sequence_kernel_name,
        step_kernel=step_kernel_name,
        return_last_states=True,
    )
    mlstm = xLSTMLarge(mlstm_config).to(device=device)

    if prefill_length > 0:
        prefill_tokens = torch.randint(0, vocab_size, (batch_size, prefill_length)).to(
            device=device
        )
    else:
        prefill_tokens = None
        batch_size = 1

    generated_tokens, state = mlstm.generate(
        prefill_tokens=prefill_tokens,
        max_length=max_length,
    )

    assert generated_tokens.shape == (
        batch_size,
        max_length,
    ), f"generated_tokens has wrong shape, got {generated_tokens.shape}"
Tests calling the generation from the model.
tests/test_generation.py