Jit lookup of precomputed function

[Binbose]

I have bases functions that I am precomputing by discretizing the definition space [0,1] and storing it and its derivatives in a dictionary. During runtime, I want to lookup the value (and the value of the gradients) using this code

@custom_jvp
def f_cached(x, i, cached_bases_dict, n_derivative=0):
      x = (x * cached_bases_dict[0].shape[-1]).astype(np.int32)
      return cached_bases_dict[n_derivative][i][x]

@f_cached.defjvp
def f_fwd(primals, tangents):
   x, i, cached_bases_dict, n_derivative = primals
   t_x = tangents

   grad = cached_bases_dict[n_derivative+1] * t_x
   return M_cached(x, i, cached_bases_dict, n_derivative=n_derivative), grad

However, in line cached_bases_dict[n_derivative][i][x] I am getting this error:

TypeError: unhashable type: 'DynamicJaxprTracer'

because x is not static. Is there a way to get around this problem? Precomputing basis functions and looking them up seems to be a common thing to do.

[YouJiacheng]

You can use an jnp.array as your cached_bases_dict