Fix MCMC?

JaxGaussianProcesses · Sep 16, 2022 · fe1cf90 · fe1cf90
1 parent 1fb92ee
commit fe1cf90
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/examples/tfp_integration.ipynb b/examples/tfp_integration.ipynb
@@ -223,7 +223,7 @@
     "def build_log_pi(target, mapper_fn):\n",
     "    def array_mll(parameter_array):\n",
     "        parameter_dict = mapper_fn([jnp.array(i) for i in parameter_array])\n",
-    "        gpx.constrain(parameter_dict, bijectors)\n",
+    "        parameter_dict = gpx.constrain(parameter_dict, bijectors)\n",
     "        return target(parameter_dict)\n",
     "\n",
     "    return array_mll\n",
@@ -277,8 +277,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#unconstrained_params = gpx.unconstrain(params, bijectors)\n",
-    "#states, log_probs = jax.jit(run_chain)(key, jnp.array(dict_to_array(unconstrained_params)))\n",
+    "unconstrained_params = gpx.unconstrain(params, bijectors)\n",
+    "states, log_probs = jax.jit(run_chain)(key, jnp.array(dict_to_array(unconstrained_params)))\n",
     "states, log_probs = jax.jit(run_chain)(key, jnp.array(dict_to_array(params)))"
    ]
   },