Numpyro: Support for bijective transforms for use in HMC

[neerajprad]

This will be important for inference on models that contain distributions with constrained support. I think PyTorch's transforms module is pretty nice, and we can follow the same pattern here.

cc. @fehiepsi

[fehiepsi]

I'll try to make an initial api for this after adding dirichlet distribution.

[neerajprad]

@fehiepsi - have you already started working on this? I think this is a blocker for working on more interesting models.

[fehiepsi]

@neerajprad I have not yet worked on this. Feel free to work on this if you want. I am doing on the hmc pr which we discussed yesterday and will work on lax.scan instead. :)

[fritzo]

See also "constraint support in minipyro" https://github.com/pyro-ppl/pyro/issues/1819
I think you should be able to:

1. fork torch.distributions.constraints
2. fork torch.distributions.transforms
3. use torch.distributions.constraint_registry's biject_to directly (or fork to avoid torch dependency)

[fehiepsi]

@fritzo , we already did this for most constraints. ^^