Method main

docs/how_to/tutorials/mix_python_and_tvm_with_pymodule.py:363–369

(
            x: R.Tensor((4, 8), "float32"),
        )

        @R.function
        def main(
            x: R.Tensor((4, 8), "float32"),
        ) -> R.Tensor((4, 8), "float32"):
            # The VM calls back into Python for these two ops
            h = R.call_py_func("layer_norm", (x,), out_sinfo=R.Tensor((4, 8), "float32"))
            out = R.call_py_func("silu", (h,), out_sinfo=R.Tensor((4, 8), "float32"))
            return out

    mod = HybridVMModule(device=tvm.cpu(0))
    x = torch.randn(4, 8)