Suggestion: alias `Transformed`(WithState) apply to __call__

[hylkedonker]

Suggestion: alias `Transformed`(WithState) apply to __call__

[tomhennigan]

Hey @hylkedonker, the con you've pointed to is the main reason we have not gone for this historically, namely that it is confusing to have a decorator that changes the signature of the function.

Haiku (like Optax, Stax and other libraries) produces a pair of functions with different signatures to the input function. You can actually treat this object as a pair if you would like:

def f(x):
  ..

f_init, f_apply = hk.transform(f)

And if you prefer to continue to reference the name f you can actually override that in the unpacking:

def f(x):
  ..

f_init, f = hk.transform(f)

params = f_init(..)
out = f(params, ..)

Again we don't suggest this in our guides since it is quite confusing for new users that the signature of f appears to have changed from the function definition.

If you really prefer stacking decorators, then you can monkey patch Haiku in place to support this functionality:

hk.Transformed.__call__ = property(fget=lambda self: self.apply)
hk.TransformedWithState.__call__ = property(fget=lambda self: self.apply)

Here is a worked example: https://colab.research.google.com/gist/tomhennigan/d9957a2233604e116d8f4ba4d9c9e3bb/add-__call__-to-transformed-withstate.ipynb

[hylkedonker]

Thanks for clarifying!