Define NNVM operator in python with TVM?

Questions
1

Hello,
I’m curious if it’s possible to define a custom NNVM operator in python via it’s lowering to a tvm schedule?

Basically some way to stick a little bit of TVM code in the middle of my NNVM graph, without having to muck around with the underlying C++.

2

Update: It’s easy enough using nnvm.top.registry.register_compute.

e.g., this bit of code:

@reg.register_compute("clip", level=15)
def compute_clip(attrs, inputs, _):
    """ Clip operator.
    """
    x = inputs[0]
    a_min = attrs.get_float("a_min")
    a_max = attrs.get_float("a_max")
    const_min = tvm.const(a_min, x.dtype)
    const_max = tvm.const(a_max, x.dtype)
    with tvm.tag_scope(topi.tag.ELEMWISE):
        x = tvm.compute(
            x.shape, lambda *i: tvm.min(x(*i), const_max), name="clipA")
        x = tvm.compute(
            x.shape, lambda *i: tvm.max(x(*i), const_min), name="clipB")
    return x