Parameter refactor v4

[slinderman]

Some pretty big changes in here...

• [x] Converting params and props to NamedTuples
• [x] Removing chex dependency
• [x] Rewriting to/from_unconstrained to use tree_map
• [x] Rewriting log_det_jacobian_constrain to use tree_map
• [x] Making all inference types NamedTuples
• [x] Making HMM smoother functions compute transition probabilities by default (kwarg specified)
• [x] Doing some general cleanup of the typing
• [x] Removing unused imports in many files
• [x] Fixing HMC notebook to use log det Jacobian
• [x] Updating tests and notebooks to make sure everything still runs

I tried a few ways of combining parameters and their properties into one object, but in the end I think this is the most transparent approach. I converted them all to NamedTuple, which seems like JAX best practice. Note that props are still mutable, which could in theory cause some problems if you jit a function that closes over ParameterProperties and then change them.

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[murphyk]

Looks great! (See comments in slack.) There are 2 failing tests, but I will merge anyway, and we can resolve this later.