Function test_tensorrt_conv3d_transpose

tests/python/relax/test_codegen_tensorrt.py:251–269

()

def test_tensorrt_conv3d_transpose():
    # Default IODHW kernel layout ([in, out, d, h, w]); output channels are weight_shape[1].
    @tvm.script.ir_module
    class ConvTranspose3d:
        @R.function
        def main(
            data: R.Tensor((2, 4, 8, 8, 8), "float32"), weight: R.Tensor((4, 2, 3, 3, 3), "float32")
        ):
            with R.dataflow():
                out = relax.op.nn.conv3d_transpose(data, weight, padding=1)
                R.output(out)
            return out

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


def test_tensorrt_int8_calibration(monkeypatch):