Tracer error for pure_callbacks

[SNMS95]

I was trying to play around with the example given in https://jax.readthedocs.io/en/latest/external-callbacks.html for working with jax.pure_callback. I made a slight modification to the example and it shows an UnexpectedTracerError.

import jax
import jax.numpy as jnp
import scipy.special
from functools import partial

@partial(jax.custom_jvp, nondiff_argnums=(2,))
def jv(v, z, problem_dict):
  v, z = jnp.asarray(v), jnp.asarray(z)
  
  # My edits
  v_scale = problem_dict['scale_v']
  z_scale = problem_dict['scale_z']
  v = v * v_scale
  z = z / z_scale
  # End of my edits

  assert jnp.issubdtype(v.dtype, jnp.integer)
  z = z.astype(jnp.result_type(float, z.dtype))
  _scipy_jv = lambda v, z: scipy.special.jv(v, z).astype(z.dtype)
  result_shape_dtype = jax.ShapeDtypeStruct(
      shape=jnp.broadcast_shapes(v.shape, z.shape),
      dtype=z.dtype)
  return jax.pure_callback(_scipy_jv, result_shape_dtype, v, z, vectorized=True)

@jv.defjvp
def _jv_jvp(problem_dict, primals, tangents):
  v, z = primals
  _, z_dot = tangents  # Note: v_dot is always 0 because v is integer.
  # Added the third argument to the jv function
  jv_minus_1, jv_plus_1 = jv(v - 1, z, problem_dict), jv(v + 1, z, problem_dict)
  djv_dz = jnp.where(v == 0, -jv_plus_1, 0.5 * (jv_minus_1 - jv_plus_1))
  return jv(v, z, problem_dict), z_dot * djv_dz

j1 = partial(jv, 1)
problem_dict = {'scale_v': 1, 'scale_z': 1.0}
print(jax.grad(jax.jit(j1))(2.0, problem_dict))

The only change is made is to introduce a new argument problem_dict but all functions involved seem pure to me.
Why does this cause the error?

The error stems at jv_minus_1, jv_plus_1 = jv(v - 1, z, problem_dict), jv(v + 1, z, problem_dict) but I couldn't find out why

[SNMS95]

Update

1. jax.jit(jax.grad) works! but jax.grad(jax.jit) is the one causing issues

[jakevdp]

The issue is that you've marked problem_dict as a nondiff_argnum, which requires it to be static. So you must also mark it as a static_argnum when using jit.

[jakevdp]

That said, static arguments to jit must be hashable, and dicts are not hashable, so instead you could keep it static using a closure:

problem_dict = {'scale_v': 1, 'scale_z': 1.0}
j1 = partial(jv, 1, problem_dict=problem_dict)
print(jax.grad(jax.jit(j1))(2.0))

[SNMS95]

Also, just for people who may not know, we shouldn't use non diff_args with Traced objects.
By modifying the JVP rule and removing it makes it work as well!