Method main

tests/python/relax/test_pytorch_integration.py:47–63

Main function demonstrating cross-function calls.

(self, x: torch.Tensor, w: torch.Tensor)

    @I.pyfunc
    def main(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
        """Main function demonstrating cross-function calls."""
        n = x.shape[0]

        # Call TIR function
        lv = self.call_tir(self.matmul, [x, w], out_sinfo=R.Tensor((n, 20), "float32"))

        # Apply ReLU
        lv1 = F.relu(lv)

        # Call packed function (will be added dynamically)
        lv2 = self.call_dps_packed("my_softmax", [lv1, 1], out_sinfo=R.Tensor((n, 20), "float32"))

        # Call Python function
        lv3 = self.my_identity_func(lv2)

        return lv3

    @T.prim_func(s_tir=True)
    def matmul(