| 46 | class ExpectedModule: |
| 47 | @R.function |
| 48 | def forward( |
| 49 | x: R.Tensor((1, 10), dtype="float32"), |
| 50 | packed_params: R.Tuple( |
| 51 | R.Tensor((20, 10), dtype="float32"), R.Tensor((20, 10), dtype="float32") |
| 52 | ), |
| 53 | ): |
| 54 | R.func_attr({"num_input": 1}) |
| 55 | with R.dataflow(): |
| 56 | linear_1_weight = packed_params[0] |
| 57 | linear_2_weight = packed_params[1] |
| 58 | matmul_1_weight = R.permute_dims(linear_1_weight) |
| 59 | matmul = R.matmul(x, matmul_1_weight) |
| 60 | matmul_2_weight = R.permute_dims(linear_2_weight) |
| 61 | matmul1 = R.matmul(x, matmul_2_weight) |
| 62 | add = R.add(matmul, matmul1) |
| 63 | gv = add |
| 64 | R.output(gv) |
| 65 | return gv |
| 66 | |
| 67 | model = TestModule(10, 20) |
| 68 | mod, _ = model.export_tvm( |