Is it desirable to have an objax.Function module to wrap functions and their vars?

[david-berthelot]

When passing a function to Jit or Grad, etc... we need to also pass the variables the function uses. While we can maintain this design pattern, we could also propose an alternative: a decorator that turns a function into a module.

Example:

# Current design
def some_function(x):
    return objax.functional.softmax(model(x, training=False))

jit_func = objax.Jit(some_function, model.vars())

Alternate design proposal (not a replacement)

@objax.Function(model.vars())
def some_function(x):
    return objax.functional.softmax(model(x, training=False))

jit_func = objax.Jit(some_function)

I like a certain elegance about this proposal but I would like more opinions on whether it is a desirable addition.