[WIP] Type hints for NN and param store#2865

kamathhrishi wants to merge 6 commits intopyro-ppl:devpyro-ppl/pyro:devfrom

kamathhrishi:devCopy head branch name to clipboard