| 889 | |
| 890 | @R.function |
| 891 | def func2( |
| 892 | x: R.Tensor((2, 3), dtype="float32"), y: R.Tensor((2, 3), dtype="float32") |
| 893 | ) -> R.Tensor((2, 3), dtype="float32"): |
| 894 | R.func_attr({"relax.force_pure": True}) |
| 895 | cls = Expected |
| 896 | storage: R.Object = R.memory.alloc_storage( |
| 897 | R.shape([24]), virtual_device_index=0, storage_scope="global", dtype="float32" |
| 898 | ) |
| 899 | alloc: R.Tensor((2, 3), dtype="float32") = R.memory.alloc_tensor( |
| 900 | storage, 0, R.shape([2, 3]), dtype="float32" |
| 901 | ) |
| 902 | _: R.Tuple() = cls.add(x, x, alloc) |
| 903 | gv1: R.Tensor((2, 3), dtype="float32") = alloc |
| 904 | alloc1: R.Tensor((2, 3), dtype="float32") = R.memory.alloc_tensor( |
| 905 | storage, 0, R.shape([2, 3]), dtype="float32" |
| 906 | ) |
| 907 | _2: R.Tuple() = cls.add(y, y, alloc1) |
| 908 | gv12: R.Tensor((2, 3), dtype="float32") = alloc1 |
| 909 | return x |
| 910 | |
| 911 | mod = relax.transform.StaticPlanBlockMemory()(Module) |
| 912 | tvm.ir.assert_structural_equal(mod, Expected) |