(
rank_arg: R.Prim(value="rank"),
fget_item: R.Callable([R.Object, R.Prim("int64")], R.Object),
)
| 41 | |
| 42 | @R.function |
| 43 | def transform_params( |
| 44 | rank_arg: R.Prim(value="rank"), |
| 45 | fget_item: R.Callable([R.Object, R.Prim("int64")], R.Object), |
| 46 | ): |
| 47 | rank = T.int64() |
| 48 | |
| 49 | A = fget_item(R.str("A"), R.prim_value(0)) |
| 50 | A = R.match_cast(A, R.Tensor([4, 4], "int32")) |
| 51 | A = R.strided_slice(A, axes=[0], begin=[rank * 2], end=[(rank + 1) * 2]) |
| 52 | |
| 53 | B = fget_item(R.str("B"), R.prim_value(1)) |
| 54 | B = R.match_cast(B, R.Tensor([2, 2], "float32")) |
| 55 | B = R.strided_slice(B, axes=[1], begin=[rank * 1], end=[(rank + 1) * 1]) |
| 56 | |
| 57 | return (A, B) |
| 58 | |
| 59 | pipeline = tvm.ir.transform.Sequential( |
| 60 | [ |
no test coverage detected
searching dependent graphs…