Skip to content

Fix documentation for samplers

Vicentini Filippo requested to merge github/fork/wdphy16/sampler_api_2 into master

Created by: wdphy16

The final one in the #1013 series (except the tests for PT samplers). Now I think the sampler API is ready for the NetKet 3.3 release.

I reviewed all docstrings and type annotations in sampler/base.py and sampler/metropolis.py, which made me confused when I read them for the first time. But I didn't touch numpy, PT, and pmap samplers.

Most type annotations for model are actually not just nn.Module, but Union[nn.module, WrappedApplyFun, HaikuWrapper, ...], because we once stated that we would support other frameworks like Haiku, although there is not much effort into that recently. For now as we're not doing strict type linting, let's just use duck typing and put nn.Module there.

Merge request reports