Function test_executable_jit_force_recompile

tests/python/runtime/test_executable.py:231–259

Test jit method with force_recompile=True.

()

def test_executable_jit_force_recompile():
    """Test jit method with force_recompile=True."""
    # Create target and build
    target = tvm.target.Target("c")
    lib = tvm.tirx.build(MyModule, target=target)

    # Create an executable
    executable = Executable(lib)

    # First jit call
    jitted_mod1 = executable.jit()

    # Second jit call without force_recompile should return the same module
    jitted_mod2 = executable.jit()
    assert jitted_mod1 is jitted_mod2

    # Third jit call with force_recompile should return a new module
    jitted_mod3 = executable.jit(force_recompile=True)
    assert jitted_mod3 is not jitted_mod1

    # Test the function works
    a = tvm.runtime.tensor(np.array([1.0] * 10, dtype="float32"))
    b = tvm.runtime.tensor(np.array([2.0] * 10, dtype="float32"))
    c = tvm.runtime.tensor(np.array([0.0] * 10, dtype="float32"))

    jitted_mod3["add"](a, b, c)

    # Check results
    tvm.testing.assert_allclose(c.numpy(), np.array([3.0] * 10, dtype="float32"))


if __name__ == "__main__":