Function tiny_runtime

tests/test_runtime.py:15–41

Create a TIDERuntime with a tiny model and random routers.

(tmp_path)

@pytest.fixture
def tiny_runtime(tmp_path):
    """Create a TIDERuntime with a tiny model and random routers."""
    model = TinyModel(num_layers=8, dim=64, vocab_size=256)
    hidden_dim = 64
    bottleneck_dim = 32

    # Create routers at checkpoint layers (3, 7)
    routers = {
        3: TokenRouter(hidden_dim, bottleneck_dim),
        7: TokenRouter(hidden_dim, bottleneck_dim),
    }

    checkpoint = RouterCheckpoint(
        routers=routers, hidden_dim=hidden_dim, bottleneck_dim=bottleneck_dim
    )
    router_path = str(tmp_path / "router.pt")
    checkpoint.save(router_path)

    config = TIDEConfig(
        checkpoint_interval=4,
        exit_threshold=0.5,
        min_layers=0,
        router_bottleneck_dim=bottleneck_dim,
    )

    runtime = TIDERuntime(model, router_path, config=config, use_cuda_kernels=False)
    return runtime, model


class TestTIDERuntimeOutput: