Use optimized graph for JAX conversion

Author: brandonwillard

pymc3/sampling_jax.py

@@ -47,7 +47,7 @@ def sample_tfp_nuts(
 
     seed = jax.random.PRNGKey(random_seed)
 
-    fgraph = aesara.graph.fg.FunctionGraph(model.free_RVs, [model.logpt])
+    fgraph = model.logp.f.maker.fgraph
     fns = jax_funcify(fgraph)
     logp_fn_jax = fns[0]