Индексирование словаря с помощью Numpy/Jax

Я пишу процедуру интерполяции и имею словарь, в котором хранятся значения функции в точках подгонки. В идеале, ключи словаря должны представлять собой 2D-массивы Numpy координат подходящей точки, np.array([x, y]), но поскольку массивы Numpy не подлежат хэшированию, они преобразуются в кортежи для ключей.

# fit_pt_coords: (n_pts, n_dims) array
# fn_vals: (n_pts,) array
def fit(fit_pt_coords, fn_vals):
    pt_map = {tuple(k): v for k, v in zip(fit_pt_coords, fn_vals)}
    ...

Далее в коде мне нужно получить значения функции, используя координаты в качестве ключей, чтобы выполнить интерполяцию. Мне бы хотелось, чтобы это было в коде @jax.jit, но значения координат имеют тип <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>, который невозможно преобразовать в кортеж. Я пробовал и другие вещи, например, создание ключа словаря как (x + y, x - y), но опять же, для этого требуются конкретные значения, и вызов .item() приводит к ConcretizationTypeError.

На данный момент я @jax.jitиспользовал весь код, который только мог, и только что оставил этот код необработанным. Однако было бы здорово, если бы я мог использовать и этот код. Есть ли какие-нибудь лучшие способы индексации словаря (или лучшие структуры данных, совместимые с Jax), которые позволили бы отбрасывать весь код? Я новичок в Jax и до сих пор не понимаю, как он работает, поэтому уверен, что должны быть лучшие способы сделать это...