Method main

tests/python/contrib/test_example_npu.py:72–83

(
        x: R.Tensor((1, 16, 32, 32), "float32"),
    )

    @R.function
    def main(
        x: R.Tensor((1, 16, 32, 32), "float32"),
    ) -> R.Tensor((1, 16, 16, 16), "float32"):
        with R.dataflow():
            # First ReLU
            y = relax.op.nn.relu(x)
            # Max pooling
            z = relax.op.nn.max_pool2d(y, pool_size=(2, 2), strides=(2, 2))
            # Second ReLU
            out = relax.op.nn.relu(z)
            R.output(out)
        return out


@tvm.script.ir_module