Function expected2

tests/python/relax/test_transform_combine_parallel_matmul.py:108–123

(
        x: R.Tensor((2, 1024, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
    )

    @R.function
    def expected2(
        x: R.Tensor((2, 1024, 640), dtype="float32"),
        y: R.Tensor((640, 640), dtype="float32"),
        y_1: R.Tensor((640, 640), dtype="float32"),
        y_2: R.Tensor((640, 640), dtype="float32"),
    ) -> R.Tensor((2, 3072, 640), dtype="float32"):
        with R.dataflow():
            lv = R.concat((y, y_1, y_2), axis=1)
            lv1 = R.matmul(x, lv, out_dtype="float32")
            lv2 = R.split(lv1, indices_or_sections=[640, 1280], axis=2)
            lv_1 = lv2[0]
            lv1_1 = lv2[1]
            lv2_1 = lv2[2]
            lv3 = R.concat((lv_1, lv1_1, lv2_1), axis=1)
            R.output(lv3)
        return lv3

    tvm.ir.assert_structural_equal(mod["main"], expected2.with_attr("global_symbol", "main"))