Function get_mod

tests/python/relax/test_codegen_coreml.py:202–210

(axis)

@requires_coreml_runtime
def test_expand_dims():
    def get_mod(axis):
        x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
        bb = relax.BlockBuilder()
        with bb.function("main", [x]):
            with bb.dataflow():
                lv0 = bb.emit(relax.op.expand_dims(x, axis=axis))
                gv = bb.emit_output(lv0)
            bb.emit_func_output(gv)
        return bb.get()

    x_data = tvm.runtime.tensor(np.random.rand(10, 10).astype("float32"), dev)
    verify(get_mod(axis=0), [x_data])