Import jax in TFP function. Fix time measurement of sampler. (#4290)

Author: twiecki

pymc3/sampling_jax.py

@@ -38,7 +38,7 @@ def sample_tfp_nuts(
     num_compute_step_size=500,
 ):
     from tensorflow_probability.substrates import jax as tfp
-
+    import jax
     model = modelcontext(model)
 
     seed = jax.random.PRNGKey(random_seed)
@@ -97,8 +97,6 @@ def get_tuned_stepsize(samples, step_size):
     tic2 = pd.Timestamp.now()
     map_seed = jax.random.split(seed, chains)
     mcmc_samples, leapfrog_num = _sample(init_state_batched, map_seed)
-    tic3 = pd.Timestamp.now()
-    print("Compilation + sampling time = ", tic3 - tic2)
 
     # map_seed = jax.random.split(seed, chains)
     # mcmc_samples = _sample(init_state_batched, map_seed)
@@ -108,10 +106,10 @@ def get_tuned_stepsize(samples, step_size):
     posterior = {k: v for k, v in zip(rv_names, mcmc_samples)}
 
     az_trace = az.from_dict(posterior=posterior)
+    tic3 = pd.Timestamp.now()
+    print("Compilation + sampling time = ", tic3 - tic2)
     return az_trace  # , leapfrog_num, tic3 - tic2
 
-    import jax
-
 
 def sample_numpyro_nuts(
     draws=1000,
@@ -169,9 +167,6 @@ def _sample(current_state, seed):
     tic2 = pd.Timestamp.now()
     map_seed = jax.random.split(seed, chains)
     mcmc_samples, leapfrogs_taken = _sample(init_state_batched, map_seed)
-    tic3 = pd.Timestamp.now()
-    print("Compilation + sampling time = ", tic3 - tic2)
-
     # map_seed = jax.random.split(seed, chains)
     # mcmc_samples = _sample(init_state_batched, map_seed)
     # tic4 = pd.Timestamp.now()
@@ -180,4 +175,6 @@ def _sample(current_state, seed):
     posterior = {k: v for k, v in zip(rv_names, mcmc_samples)}
 
     az_trace = az.from_dict(posterior=posterior)
+    tic3 = pd.Timestamp.now()
+    print("Compilation + sampling time = ", tic3 - tic2)
     return az_trace  # , leapfrogs_taken, tic3 - tic2