Numpyro

[thomaspinder]

This PR ports the existing Distrax functionality into NumPyro according to #121

Pull request type

• Bugfix
• Feature
• Code style update (formatting, renaming)
• [ x ] Refactoring (no functional changes, no api changes)
• Build related changes
• Documentation content changes
• Other (please describe):

Issue Number: #121

[codecov]

Codecov Report

❗ No coverage uploaded for pull request base (v0.5_update@bb26a8e). Click here to learn what that means.
The diff coverage is n/a.

@@              Coverage Diff               @@
##             v0.5_update     #122   +/-   ##
==============================================
  Coverage               ?   99.23%           
==============================================
  Files                  ?       14           
  Lines                  ?     1172           
  Branches               ?        0           
==============================================
  Hits                   ?     1163           
  Misses                 ?        9           
  Partials               ?        0           

Flag	Coverage Δ
unittests	99.23% <0.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[daniel-dodd]

Great PR. Main comment is about Dict vs dict consistency. Issues with docs have already been address at my end, so ignore these comments - they are mainly for myself when I come to rebase.

[daniel-dodd]

Dictionary typing needs to be consistent - your changes have replaced some methods with Dict to dict, while others are unchanged . I opt to us Dict as it more general, though happy to switch to dict if there's good reason - though all methods need to be changed throughout the library.

[thomaspinder]

These have been reolved now.