Function load_module

tests/python/relax/test_runtime_sampling_flashinfer.py:33–43

(name: str, static_modules: list[tvm.runtime.Module])

@pytest.mark.skip(reason="Requires FlashInfer enabled and proper setup")
def test_sampling():
    def load_module(name: str, static_modules: list[tvm.runtime.Module]):
        assert len(static_modules) > 0
        if len(static_modules) == 1:
            return static_modules[0]
        static_mod = static_modules[0]
        for mod in static_modules[1:]:
            static_mod.import_module(mod)
        temp = utils.tempdir()
        mod_path = temp.relpath(f"{name}.so")
        static_mod.export_library(mod_path)
        return tvm.runtime.load_module(mod_path)

    # Test configuration
    batch_size = 10