def annotated():
x = relay.var("x", relay.ty.TensorType((3, 3, 4), "float32"))
split = relay.op.split(x, 3)
split = split.astuple()
split = relay.annotation.on_device(split, gpu_dev)
split = relay.TupleWrapper(split, 3)
sub = split[0] - split[1]
func = relay.Function(relay.analysis.free_vars(sub), sub)
func = run_opt_pass(func, transform.RewriteAnnotatedOps(cpu_dev.device_type))
return func
def annotated():
x = relay.var("x", relay.ty.TensorType((3, 3, 4), "float32"))
split = relay.op.split(x, 3)
split = split.astuple()
split = relay.annotation.on_device(split, gpu_ctx)
split = relay.TupleWrapper(split, 3)
sub = split[0] - split[1]
func = relay.Function(relay.ir_pass.free_vars(sub), sub)
func = relay.ir_pass.rewrite_annotated_ops(func, cpu_ctx.device_type)
return func
