Function test_tensorrt_conv1d

tests/python/relax/test_codegen_tensorrt.py:137–152

()

def test_tensorrt_conv1d():
    # Regression test: explicit-batch (batch > 1) 1D convolution. The pre-TRT10 converter assumed an
    # implicit batch dimension and dropped the spatial dimension under explicit batch.
    @tvm.script.ir_module
    class Conv1d:
        @R.function
        def main(data: R.Tensor((2, 8, 16), "float32"), weight: R.Tensor((4, 8, 3), "float32")):
            with R.dataflow():
                out = relax.op.nn.conv1d(data, weight, padding=1)
                R.output(out)
            return out

    data = np.random.randn(2, 8, 16).astype("float32")
    weight = np.random.randn(4, 8, 3).astype("float32")
    patterns = [("tensorrt.nn.conv1d", is_op("relax.nn.conv1d")(wildcard(), wildcard()))]
    _offload_and_compare(Conv1d, {"weight": weight}, patterns, data)


def test_tensorrt_max_pool2d():