[Op] Print

[comaniac]

Description

This PR introduces a utility op op.Print that prints intermediate results for traced modules. A unit test and tutorial are also provided.

Checklist

• PR's title starts with a category (e.g. [Bugfix], [Model], [Tutorial], etc)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage
• Code is well-documented

cc @chhzh123 @zarzen

[chhzh123]

I'm actually thinking some non-intrusive ways to print the intermediate output. Otherwise, users still need to modify the model itself.

[chhzh123]

For example, can we add a .print() primitive to do this? It should act like .annotate that only adds a hook or insert a module for printing the results.

[chhzh123]

self.linear

[comaniac]

For example, can we add a .print() primitive to do this? It should act like .annotate that only adds a hook or insert a module for printing the results.

This also creates some limitations. For example, you cannot print intermediate results inside a module but just its final output. We could maybe design a primitive only for traced GraphModule, but in this case users still have to use .find to identify the position for printing. We could think more on this mechanism and see if we could further improve it in the future.

[chhzh123]

I can think of two ways for now:

1. Attach print information to submodules and print out the results when executing these submodules, but this does not support arbitrary print point as you mentioned.
2. Leverage fx.Interpreter. We can either print the output for each op, or provide an interface for users to determine which op's output should be print out.

[chhzh123]

@comaniac We can merge for now and improve this feature later.

[comaniac]

Sure please go ahead and merge then.

[chhzh123]

Thanks