Support nn.Modules in @function, make_funsor, and make_op - Githubissues

pyro-ppl / funsor

Functional tensors for probabilistic programming

https://funsor.pyro.ai

Apache License 2.0

236 stars 20 forks source link

Support nn.Modules in @function, make_funsor, and make_op #507

Closed fritzo closed 3 years ago

fritzo commented 3 years ago

This adds support for torch.nn.Module instances in @function, make_funsor(), and make_op.

This also fixes a couple unrelated docs errors:

addes docstring to examples/forward_backward.py
changes scope of Op.arity cvar

Tested

[x] refactored VAE example to test @function (this example will later be changed to use make_op())
[x] unit test for make_funsor()
[x] unit test for make_op()