Function test_batch_flatten

tests/python/relax/test_codegen_coreml.py:233–244

()

@requires_coreml_runtime
def test_batch_flatten():
    x = relax.Var("x", relax.TensorStructInfo([10, 10, 10], "float32"))
    bb = relax.BlockBuilder()
    with bb.function("main", [x]):
        with bb.dataflow():
            lv0 = bb.emit(relax.op.nn.batch_flatten(x))
            gv = bb.emit_output(lv0)
        bb.emit_func_output(gv)
    mod = bb.get()

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


@requires_coreml_runtime