Function _offload_and_compare

tests/python/relax/test_codegen_tensorrt.py:117–134

Offload a single-op module to TensorRT and compare against the LLVM reference. Each module here contains a single instance of the op under test, which both exercises the individual converter and avoids the structurally-identical-composite deduplication that would otherwise collapse repe

(mod, params_np, patterns, data_np, rtol=1e-2, atol=1e-2)

def _offload_and_compare(mod, params_np, patterns, data_np, rtol=1e-2, atol=1e-2):
    """Offload a single-op module to TensorRT and compare against the LLVM reference.

    Each module here contains a single instance of the op under test, which both exercises the
    individual converter and avoids the structurally-identical-composite deduplication that would
    otherwise collapse repeated ops.
    """
    ref = build_and_run(mod, [data_np, *params_np.values()], "llvm", legalize=True)
    offloaded = tvm.transform.Sequential(
        [
            relax.transform.BindParams("main", params_np),
            relax.transform.FuseOpsByPattern(patterns),
            relax.transform.MergeCompositeFunctions(),
            relax.transform.RunCodegen(),
        ]
    )(mod)
    out = build_and_run(offloaded, [data_np], "cuda")
    tvm.testing.assert_allclose(out, ref, rtol=rtol, atol=atol)


def test_tensorrt_conv1d():