Class HybridVMModule

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

    @I.ir_module
    class HybridVMModule(BasePyModule):
        @I.pyfunc
        def silu(self, x):
            """SiLU/Swish activation — using Python as fallback."""
            return torch.sigmoid(x) * x

        @I.pyfunc
        def layer_norm(self, x):
            """LayerNorm — another Python fallback."""
            return F.layer_norm(x, x.shape[-1:])

        @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)