Add missing symbol to test

test/test_jax.py

@@ -11,5 +11,5 @@ cosx = 1.0 * sympy.cos(x) + y
 key = random.PRNGKey(0)
 X = random.normal(key, (1000, 2))
 true = 1.0 * jnp.cos(X[:, 0]) + X[:, 1]
-f, params = sympy2jax(cosx, [x])
+f, params = sympy2jax(cosx, [x, y, z])
 assert jnp.all(jnp.isclose(f(X, params), true)).item()