@@ -0,0 +1,379 @@
+import jax.numpy as np%0Afrom jax import jit, grad, jvp, make_jaxpr%0Afrom jax import primitive%0A%0A@primitive  # try commenting me out%0Adef library_fun(x):%0A  return np.sin(np.cos(x))%0A%0Adef user_fun(x):%0A  y = library_fun(x)%0A  z = np.tanh(x)%0A  return y + z%0A%0A%0Aprint user_fun(3.)%0Aprint jit(user_fun)(3.)%0Aprint grad(user_fun)(3.)%0Aprint grad(jit(user_fun))(3.)%0A%0Aprint make_jaxpr(user_fun)(3.)%0A

add example script