make_module: First version

[z-a-f]

Stack from ghstack:

• make_module: First version #23288 make_module: First version

Differential Revision: D16455390

[z-a-f]

Example:

mod = make_module(torch.add)
mod_quant = make_module(torch.add, quantized=True)

a = torch.tensor(1, dtype=torch.float)
qa = torch.quantize_linear(a, 1.0, 127, torch.quint8)
b = torch.tensor(2, dtype=torch.float)
qb = torch.quantize_linear(b, 1.0, 127, torch.quint8)

print("Types\t\t:", type(mod), type(mod_quant))
print("__repr__\t:", mod, mod_quant)
print("forward call\t:", mod(a, b))
print("forward call\t:", mod_quant(qa, qb, 1.0, 0))

Prints:

Types		: <class '__main__._WrapModule'> <class '__main__._WrapModule'>
__repr__	: Add() Add_quantized()
forward call	: tensor(3.)
forward call	: tensor(3., size=(), dtype=torch.quint8, scale=1.0, zero_point=0)

I can make the returned type to be <class '__main__.Add'> and <class '__main__.Add_quantized'>, but that would require slightly less readable version of the code :)

[raghuramank100]

Please see suggestions

[raghuramank100]

First draft, will add more comments later.

[dzhulgakov]

My biggest comment is that we need to figure out how to make it work with TorchScript (as you want the resulting model to be scriptable and runnable without python). Given that codegen + script to codegen might be an more robust approach (if you have a closed set of modules to cover which appears to be the case)

[raghuramank100]

Please see comments

[jerryzh168]

LGTM, we can change to the base class solution in future prs.

[facebook-github-bot]

@zafartahirov merged this pull request in 9c549df.