[Tracking Issue][UX][Unity] Private annotations in TVMScipt

[tqchen]

Global symbol right now needs to be manually attached, for better UX, it is desirable to move to the following model

class IRModule:
    # implies a private function
    @R.function(private=True)
    def func1():
        pass

    # implies a function with global_symbol=func2, and we do not need to fill in the attr field.
    @R.function
    def func2():
         pass

cc @quic-sanirudh

[tqchen]

cc @yongwww @Hzfengsy @cyx-6

[slyubomirsky]

Implementing this kind of syntactic sugar would also be useful for purity tracking, FYI. Unfortunately I'm not sure Python permits using @R.function to be used with no arguments while also allowing for optional parameters.

Example:

def dec(f, label=None):
    def wrapper(*args, **kwargs):
        if label is None:
            print("No label")
        else:
            print(label)
        return f(*args, **kwargs)
    return wrapper

# this works as expected
@dec
def g(x):
    return 1

# this will not work and will complain about lacking a positional argument
# @dec(label="hi")
# def h(x):
#    return 2

This variant works but is a little cumbersome:

def dd(label=None):
     def dec(f):
         def wrapper(*args, **kwargs):
             if label is None:
                 print("No label")
             else:
                 print(label)
             return f(*args, **kwargs)
         return wrapper
     return dec

# has to be called like so...
@(dd())
def g(x):
    return x

@(dd(label="hi"))
def h(x):
    return x

Not sure if there's any way to avoid the nasty inner call 😅 (if you leave off the extra parens, the function will be treated as the label argument, so the results will not be good).

Would we be able to override the parsing of the decorator in the first case? If it parses normally (the error I get is a TypeError), then our custom parser can just have special handling for it.

Edit: ast.parse does succeed on the first example if I uncomment the second function, so handling the case in the parser is an option if the TVMScript parser never attempts to execute the decorator (I think this is a bit different from its current implementation). It's not impossible, but would require some revisions to the current parser.

[tqchen]

i think some code along the lines below would work. The key is to make f also optional argument, and return the wrapper or wrapped function depending on whether it presents

def dec(f=None, label=None):
     def wrap_function(f):
        def wrapper(*args, **kwargs):
            ...
        return functools.wraps(f, wrapper)
           
    if f is not None:
         return wrap_function(f)
    else:
         return wrap_function

[slyubomirsky]

Wow, that is a very clever solution and it indeed works. I will implement this solution for purity annotations and we can also use it for privacy too.

[slyubomirsky]

Implemented