Numpyro: Support AutoDelta

[fehiepsi]

AutoDelta is used to do MAP inference in Pyro. It would be nice to have it in NumPyro too.

References

• Pyro AutoDelta implementation
• AutoNormal implementation
• autoguide's tests

[xidulu]

@fehiepsi
I would love to implement AutoDelta for NumPyro as I need to perform some experiments related to MAP estimation.
But it seems that someone has already been working on it on their own forked version https://github.com/pyro-ppl/numpyro/issues/792#issuecomment-721637496
And their PR (https://github.com/pyro-ppl/numpyro/pull/649) seems to get blocked.
Is it possible to let me port part of their implementation into the master branch? (as well as adding some test cases)

[fehiepsi]

Hi @xidulu, please go ahead. The SteinVI PR ideas have been adopted and refactored out into smaller PRs: #667 (lift handler), #804 (prng_key primitive), #799 (svi run). It would be nice to have a separate PR for AutoDelta too. I think the implementation will be a bit different from the one in SteinVI PR, e.g. no need to have a method find_params (this is implemented in AutoGuide._setup_prototype), we can reuse create_plates logics in AutoNormal.

[fehiepsi]

Resolved in #816. Thanks, @xidulu!