jax

0d053f0e - temporarily revert #1658 due to TFP test failures

temporarily revert #1658 due to TFP test failures This commit unfortunately un-fixes #1571, but only until we sort out why a TF Probvability test started failing.