Support for mutable params

[fehiepsi]

Close #990.

This PR introduces an additional mutable_state field in SVIState, which stores parameters marked by infer={"mutable": True}. Those parameters will be skipped in the optimizer and it is the job of the model/guide to update them, e.g.

def model(...):
    p = numpyro.param("p", {"a": 1}, infer={"mutable": True}) 
    p["a"] = 10  # update p

• Support mutable params in haiku
• Add new primitive numpyro.mutable
• Add a few assertations (e.g. params is None == mutable is None)
• Consider to replicate methods with new input/output signatures

[fritzo]

Hi @fehiepsi, could you provide some documentation somewhere (maybe in the numpyro.param docstring?) to explain what mutable params are and why they might be useful? I'm struggling to imagine a use case from the interface alone.

[fritzo]

LGTM, thanks for your patience and for explaining the code!

[fritzo]

LGTM.

I hav reviewed general design and skimmed code. I have not meticulously reviewed all logic (e.g. destructured assignment) and assume tests will catch type errors.

[fehiepsi]

Thanks for reviewing, @fritzo! Right now we don't have a good set of tests but we can exercise this pattern in PR #1056.