| 76 | |
| 77 | |
| 78 | def test_get_valid_counts(): |
| 79 | @R.function |
| 80 | def foo( |
| 81 | data: R.Tensor((10, 5, 6), "float32"), |
| 82 | ) -> R.Tuple( |
| 83 | R.Tensor((10,), "int32"), |
| 84 | R.Tensor((10, 5, 6), "float32"), |
| 85 | R.Tensor((10, 5), "int32"), |
| 86 | ): |
| 87 | gv: R.Tuple( |
| 88 | R.Tensor((10,), "int32"), |
| 89 | R.Tensor((10, 5, 6), "float32"), |
| 90 | R.Tensor((10, 5), "int32"), |
| 91 | ) = R.vision.get_valid_counts(data, score_threshold=0.5, id_index=0, score_index=1) |
| 92 | return gv |
| 93 | |
| 94 | data = relax.Var("data", R.Tensor((10, 5, 6), "float32")) |
| 95 | |
| 96 | bb = relax.BlockBuilder() |
| 97 | with bb.function("foo", [data]): |
| 98 | gv = bb.emit( |
| 99 | relax.op.vision.get_valid_counts(data, score_threshold=0.5, id_index=0, score_index=1) |
| 100 | ) |
| 101 | bb.emit_func_output(gv) |
| 102 | |
| 103 | _check(foo, bb.get()["foo"]) |
| 104 | |
| 105 | |
| 106 | def test_non_max_suppression_return_indices(): |