| 35 | |
| 36 | |
| 37 | def test_sort(): |
| 38 | @R.function |
| 39 | def foo(x: R.Tensor((2, 3), "int32")) -> R.Tuple( |
| 40 | R.Tensor((2, 2), dtype="int32"), R.Tensor((2, 2), dtype="int32") |
| 41 | ): |
| 42 | lv0 = R.sort(x, axis=1) |
| 43 | lv1 = R.argsort(lv0) |
| 44 | r = R.topk(lv1, axis=1, k=2) |
| 45 | return r |
| 46 | |
| 47 | x = relax.Var("x", R.Tensor((2, 3), "int32")) |
| 48 | bb = relax.BlockBuilder() |
| 49 | with bb.function("foo", (x,)): |
| 50 | lv0 = bb.emit(relax.op.sort(x, axis=1)) |
| 51 | lv1 = bb.emit(relax.op.argsort(lv0)) |
| 52 | r = bb.emit(relax.op.topk(lv1, axis=1, k=2)) |
| 53 | bb.emit_func_output(r) |
| 54 | |
| 55 | _check(foo, bb.get()["foo"]) |
| 56 | |
| 57 | |
| 58 | if __name__ == "__main__": |