Fix interaction between PyroParam and torch.func.grad

[eb8680]

Addresses a bug described downstream in BasisResearch/chirho#393

This PR adds a fix for compatibility of PyroModules and PyroParams with torch.func.grad and the other functional automatic differentiation transforms in torch.func. The fix is basically to replace each pyro.param statement or other interaction with the parameter store with a dummy version that does not store and retrieve a parameter tensor from a nonlocal state (which is invisible to the tracing machinery in torch.func).

Without this fix, gradient computations in torch.func.grad do not propagate to the unconstrained parameters behind constrained PyroParams even when using pyro.settings.set(module_local_param=True) and are always zero. After this fix, the functional AD system in torch.func behaves correctly with AutoGuides and other PyroModules when module_local_param=True, though it is still fundamentally incompatible with the global parameter store state when module_local_param=False.

Tested:

• Added simple regression test that fails without the fix in this PR

[eb8680]

I think I'm missing some edge cases. Marking this WIP while I strengthen the test and see if it still works...

[eb8680]

Addressed the rest of the edge cases, should be ready for review

[ordabayevy]

Looks great. I just have a question about the effectful.