Fix issue 1465: fix jit in example (jax-ml#1473)

* fix jit in example * Avoid using static_argnums on a keyword argument; use a positional argument and a wrapper function for now. Co-authored-by: Peter Hawkins <[email protected]>
spundun · Apr 1, 2020 · 651316f · 651316f
1 parent 86a4073
commit 651316f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/examples/advi.py b/examples/advi.py
@@ -54,7 +54,7 @@ def batch_elbo(logprob, rng, params, num_samples):
 # ========= Helper function for plotting. =========
 
 @partial(jit, static_argnums=(0, 1, 2, 4))
-def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
+def _mesh_eval(func, x_limits, y_limits, params, num_ticks):
     # Evaluate func on a 2D grid defined by x_limits and y_limits.
     x = np.linspace(*x_limits, num=num_ticks)
     y = np.linspace(*y_limits, num=num_ticks)
@@ -63,6 +63,8 @@ def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
     zs = vmap(func, in_axes=(0, None))(xy_vec, params)
     return X, Y, zs.reshape(X.shape)
 
+def mesh_eval(func, x_limits, y_limits, params, num_ticks=101):
+    return _mesh_eval(func, x_limits, y_limits, params, num_ticks)
 
 # ========= Define an intractable unnormalized density =========