| 969 | |
| 970 | @R.function |
| 971 | def fused( |
| 972 | x: R.Tensor((1, 1, 32, 128), dtype="float32"), |
| 973 | y: R.Tensor((2048, 128), dtype="float32"), |
| 974 | len: R.Shape(["m"]), |
| 975 | ) -> R.Tensor((1, 1, 32, 128), dtype="float32"): |
| 976 | R.func_attr({"Primitive": True}) |
| 977 | m = T.int64() |
| 978 | cls = Before |
| 979 | with R.dataflow(): |
| 980 | lv1 = R.emit_te(topi.add, x, x) |
| 981 | gv = R.call_tir( |
| 982 | cls.foo, |
| 983 | [lv1, y], |
| 984 | out_sinfo=R.Tensor((1, 1, 32, 128), dtype="float32"), |
| 985 | tir_vars=R.shape([m]), |
| 986 | ) |
| 987 | R.output(gv) |
| 988 | return gv |
| 989 | |
| 990 | @R.function |
| 991 | def main( |