| 199 | |
| 200 | @requires_coreml_runtime |
| 201 | def test_expand_dims(): |
| 202 | def get_mod(axis): |
| 203 | x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32")) |
| 204 | bb = relax.BlockBuilder() |
| 205 | with bb.function("main", [x]): |
| 206 | with bb.dataflow(): |
| 207 | lv0 = bb.emit(relax.op.expand_dims(x, axis=axis)) |
| 208 | gv = bb.emit_output(lv0) |
| 209 | bb.emit_func_output(gv) |
| 210 | return bb.get() |
| 211 | |
| 212 | x_data = tvm.runtime.tensor(np.random.rand(10, 10).astype("float32"), dev) |
| 213 | verify(get_mod(axis=0), [x_data]) |
| 214 | verify(get_mod(axis=1), [x_data]) |
| 215 | |
| 216 | |
| 217 | @requires_coreml_runtime |